Skip to content

Commit 4152d07

Browse files
DiFanruiclaude
andcommitted
[2026春季][T1-1-4] Implement roll, column_stack, mode, meshgrid, cartesian_prod
Five data-rearrangement operators implemented using ninetoothed output-driven gather pattern with ntl.load manual pointer arithmetic. - roll: permute + ntl.load gather per dimension (1D-5D) - column_stack: identity kernel with non-contiguous output views (0D-ND) - meshgrid: reuses cartesian_prod kernel (0D, list arg, xy/ij) - cartesian_prod: ntl.load + integer division/modulo per column - mode: O(K^2) with ntl.sum + ntl.where + mask (K padded to pow2) All 89 tests pass. Semantic verification for mode tie-breaking. Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 6bc90d5 commit 4152d07

20 files changed

Lines changed: 1186 additions & 0 deletions

HONOR_CODE.md

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# 2026 春季启元人工智能大赛诚信守则(Honor Code)
2+
3+
4+
本人作为 2026 春季启元人工智能大赛(以下简称"比赛")的参赛选手,郑重承诺严格遵守比赛规则及本诚信守则,秉持诚信、公正、廉洁的参赛原则,自觉维护比赛的公平性与严肃性。本人充分理解并认可,违反本准则将导致参赛资格被取消、比赛成绩作废等相应后果,且愿意承担由此产生的一切责任。
5+
6+
## 一、参赛诚信承诺
7+
8+
1. 本人保证所提交的赛题PR(Pull Request)中包含的算子实现代码及相关文档,均为本人(及参赛团队,如为团队参赛)在比赛期间独立完成或在明确标注参考来源的基础上进行开发,不存在任何欺诈、抄袭、作弊行为。
9+
10+
2. 本人承诺主动、全面、真实地披露赛题实现过程中所有参考的外部资源,尤其是开源代码资源,不隐瞒任何可能影响比赛公平性的信息。
11+
12+
3. 本人保证不采用任何不正当手段获取比赛优势,包括但不限于窃取其他参赛选手的代码成果、利用非比赛允许的工具或技术、与他人串通作弊等。
13+
14+
## 二、参考资源说明
15+
16+
本人确认已按比赛要求,将本次赛题实现过程中涉及的参考资源信息单独撰写至`REFERENCE.md`文件中,该文件将与本诚信守则一同作为PR附件提交。`REFERENCE.md`需根据实际参考情况,按以下要求完整填写,信息不完整或虚假填写将视为违反本准则:
17+
18+
**情况2:有参考外部开源代码及相关资源**
19+
20+
1. **参考开源项目/资源名称**:InfiniTensor/ninetoothed
21+
22+
2. **参考资源链接**https://github.com/InfiniTensor/ninetoothed
23+
24+
3. **参考的具体内容**:ninetoothed 框架的 Python API(`Tensor``Symbol``arrangement``application``premake` 函数签名和用法);`ntl.load` 手工指针加载;`ntl.sum` 规约求和;`ntl.where` 三元选择;`ntl.cast` 类型转换;`offsets()``data_ptr()``stride()``shape` 等元数据访问;`element_wise.arrangement``flatten().tile()` 模式。
25+
26+
4. **本人对参考内容的修改与优化说明**:所有算子均基于 ninetoothed 框架的公开 API 独立开发,未复制框架内部实现。五个算子采用 output-driven gather 模式,通过手工指针算术实现数据重排布,克服了 ninetoothed 0.25.0 arrangement 层不支持 split/concat/roll 等原语的限制。mode 算子通过 O(K²) 算法 + K_orig/K_tile mask 方案在应用层实现了规约,绕过了 arrangement 层 expand 的 bug。
27+
28+
5. **开源协议类型**:Apache 2.0
29+
30+
6. **其他补充信息**:开发过程中使用 Codex(OpenAI)进行架构咨询和代码审查,所有最终代码均由本人编写和验证。
31+
32+
## 三、禁止行为确认
33+
34+
本人明确知晓并承诺避免以下违反比赛公平性的行为,若存在以下任一情况,自愿接受比赛组委会的相应处罚:
35+
36+
1. 未经授权复制、抄袭他人(包括其他参赛选手、开源项目、商业代码)的代码、算法或技术方案,且未进行明确标注;
37+
2. 隐瞒或虚假披露参考资源信息,包括遗漏重要参考来源、伪造参考内容说明等;
38+
3. 与其他参赛选手或第三方串通,进行代码共享、成果交换等违规协作;
39+
4. 利用比赛平台漏洞、技术缺陷或非比赛允许的工具获取不正当利益;
40+
5. 伪造比赛相关证明材料、提交虚假信息;
41+
6. 其他违反比赛规则及公序良俗的不诚信行为。
42+
43+
## 四、责任与确认
44+
45+
1. 本人充分理解,比赛组委会将对所有提交的PR进行代码溯源、参考信息核查等公平性审查,若发现本人存在违反本准则的行为,有权随时取消本人的参赛资格、作废比赛成绩,情节严重的将在比赛相关平台进行公示。
46+
47+
2. 若因本人违反本准则导致比赛争议或第三方权益受损(如开源协议侵权等),本人将独立承担全部法律责任及相关损失,与比赛组委会无关。
48+
49+
3. 本人确认已仔细阅读并完全理解本诚信守则的全部内容,自愿签署本准则,接受比赛组委会的监督与审查。
50+
51+
## 五、签署信息
52+
53+
参赛选手姓名:**狄凡瑞**
54+
55+
签署日期:**2026年5月19日**

REFERENCE.md

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# REFERENCE.md — T1-1-4 赛题参考资源声明
2+
3+
本次赛题提交属于**情况 2**:有参考外部开源代码及相关资源。以下按比赛要求逐项列出每个参考资源的信息。
4+
5+
---
6+
7+
## 参考 1:PyTorch 开源项目
8+
9+
1. **参考开源项目/资源名称**:PyTorch
10+
11+
2. **参考资源链接**https://github.com/pytorch/pytorch
12+
13+
3. **参考的具体内容**
14+
- `torch.roll``torch.column_stack``torch.mode``torch.meshgrid``torch.cartesian_prod` 的 Python API 签名和语义规范(参数名称、默认值、返回值形状、dtype 规则、错误处理行为)。
15+
- `torch.mode` 的 CUDA kernel 源码(`aten/src/ATen/native/cuda/SummaryOps.cu`),确认 PyTorch 2.12 CUDA 实现采用排序 + BlockReduce 的并行拓扑,平局行为由并行归约树决定,官方文档明确标注为"undefined (but deterministic)"。九齿 kernel 采用确定性策略(最小值 + 最大下标),测试改为语义验证(验证返回值频次等于最大频次)。
16+
- `torch.column_stack` 的 C++ 实现(`aten/src/ATen/native/TensorShape.cpp`),确认 0D/1D 输入先 reshape 为 `(numel, 1)` 再沿 dim=1 拼接的语义。
17+
- `torch.meshgrid` 接受 0D 张量(视为长度 1)和单 list/tuple 参数的接口约定。
18+
19+
4. **本人对参考内容的修改与优化说明**
20+
- 所有算子的数学语义(roll 的环形平移、meshgrid/cartesian_prod 的索引映射、column_stack 的拼接逻辑、mode 的频次统计)通过输出驱动 gather 模式在九齿 application 函数内重新实现,不调用 PyTorch 参考算子完成核心计算。
21+
- roll 采用 permute + ntl.load 手工指针寻址,支持 1D-5D;浮点取模用 `(offsets - shift + size) % size` 避免 Triton 负数余数问题。
22+
- column_stack 的核心创新是利用九齿 arrangement 的 stride-aware `flatten().tile()` 直接写入非连续输出切片,无需临时缓冲区。
23+
- meshgrid 复用了 cartesian_prod 的整除取模 kernel,两者共享 `output_k[flat] = input_k[(flat // repeat_after[k]) % Sk]` 的索引公式。
24+
- mode 通过 O(K²) 逐行统计算法 + K_orig/K_tile mask 方案在应用层实现了规约,绕过了 arrangement 层 expand 的 0.25.0 bug。
25+
26+
5. **开源协议类型**:PyTorch 使用 BSD-3-Clause License。
27+
28+
6. **其他补充说明**:无。
29+
30+
---
31+
32+
## 参考 2:Ninetoothed(九齿)框架
33+
34+
1. **参考开源项目/资源名称**:ninetoothed(九齿 GPU 算子代码生成框架)
35+
36+
2. **参考资源链接**https://github.com/InfiniTensor/ninetoothed
37+
38+
3. **参考的具体内容**
39+
- `ninetoothed/generation.py`:code generator 的 `visit_Call` 方法(第 206-262 行)和 `visit_Subscript` 方法(第 268-294 行),确认 `tensor.data_ptr()``tensor.offsets(dim)``tensor.stride(dim)``tensor.source.shape[dim]` 在 application 函数中的可用性和约束。确认 `stride(dim)` 要求 dim 为正整数字面量(`stride(-1)` 存在代码生成 bug)。
40+
- `ninetoothed/tensor.py`:Tensor 类的 `tile``expand``flatten``permute``unsqueeze``squeeze``__getitem__` 方法在 arrangement 层的语义和索引映射实现。确认 arrangement 层无 split/concat/roll/scatter 等数据重排布原语。
41+
- `ninetoothed/symbol.py`:Symbol 类的算术运算符重载(`__add__``__sub__``__mul__``__floordiv__``__mod__`),确认 offsets 表达式可参与算术运算。
42+
- `ninetoothed/language.py``ntl.*` 操作通过 AST `call``attribute` 函数映射到 `triton.language.*` 的机制。
43+
- 现有算子参考:`src/ntops/kernels/element_wise.py``arrangement``flatten().tile()` 模式)和 `src/ntops/kernels/silu.py` 的 premake 结构(`functools.partial` + `Tensor` 构造 + `_cached_make` 调用链)。
44+
45+
4. **本人对参考内容的修改与优化说明**
46+
- 系统性地探针了九齿 0.25.0 的能力边界,发现 `ntl.load(ptr)` 手工指针路径可用(核心突破)、`input.source[dynamic_idx]` 下标存在 IndexError bug、`arrangement` 层的 `unsqueeze + expand` 只广播第一个元素(bug)、`break` 语句不支持等限制。
47+
- 针对 0.25.0 的限制,所有数据重排布算子采用 output-driven gather 模式,将索引映射逻辑从 arrangement 层下沉到 application 层的手工指针算术中。
48+
- mode 算子通过 `ntl.where` + `ntl.sum` + `ntl.cast` 的组合克服了双重循环变量类型不匹配的问题,通过 `valid = ntl.cast(offsets(1) < K_orig, ntl.int32)` mask 排除 padding lane。
49+
- 阅读框架源码以确认 API 边界,属于对开发工具的正常理解,不涉及对九齿框架本身的修改。
50+
51+
5. **开源协议类型**:随 ntops 仓库分发,ntops 使用 Apache 2.0 License。
52+
53+
6. **其他补充说明**:九齿是本赛题指定的开发框架,使用其为正常参赛行为。
54+
55+
---
56+
57+
## 参考 3:Triton 语言
58+
59+
1. **参考开源项目/资源名称**:Triton
60+
61+
2. **参考资源链接**https://github.com/triton-lang/triton
62+
63+
3. **参考的具体内容**
64+
- `triton.language` 的 API:`load``sum``where``cast``arange``program_id` 等原语的语义和参数约定。
65+
- 确认 `tl.load` 可以接受手工构造的指针(非仅 arrangement 生成的指针),支持 scalar pointer + manual offset 计算。
66+
- 确认 `tl.sum` 可以规约向量到标量,支持在 application 函数中实现跨维度的计数统计。
67+
- 确认 `tl.arange` 要求 range 大小为 2 的幂(导致 mode 的 K 需要补齐到 next_power_of_2)。
68+
69+
4. **本人对参考内容的修改与优化说明**
70+
- 九齿 DSL 层不直接暴露所有 Triton 原语(如 `load``sum``where``cast`),通过阅读 `generation.py` 确认 code generator 会将 `ninetoothed.language.X` 无条件映射为 `triton.language.X`,从而在 application 函数中直接使用这些原语。
71+
- 发现了 Triton `%` 对负数产生负余数的问题,roll kernel 采用 `(offsets - shift + size) % size` 保证结果非负。
72+
- 这是对开发工具链的正常理解,不涉及对 Triton 本身的修改。
73+
74+
5. **开源协议类型**:Triton 使用 MIT License。
75+
76+
6. **其他补充说明**:Triton 是九齿的底层编译后端,使用其为正常参赛行为。
77+
78+
---
79+
80+
## 参考 4:Codex(OpenAI)开发辅助
81+
82+
1. **参考资源名称**:Codex(OpenAI 编程助手)
83+
84+
2. **参考资源链接**:通过 Claude Code IDE 集成访问
85+
86+
3. **参考的具体内容**
87+
- 提供了架构设计建议:确认 output-driven gather 模式在九齿中可行,建议使用 `ntl.load` 手工指针替代 `source[idx]` 下标(因 0.25.0 bug)。
88+
- 提供了 mode 的实现方向:O(K²) 逐行统计算法、K_orig + K_tile + mask 方案、`ntl.sum` + `ntl.where` 的用法。
89+
- 审查了代码质量和 API 边界覆盖(0D 输入、单 list 参数、非 1D 报错等),指出了参数校验和测试覆盖的遗漏。
90+
- 审查了 REFERENCE.md 和 HONOR_CODE.md 的格式和完整性。
91+
92+
4. **本人对参考内容的修改与优化说明**
93+
- 所有最终代码由本人编写、调试和验证。Codex 提供的是方向性建议,具体实现(如 ndim 特化的 application 函数、安全取模公式、mask 的 `* valid` 而非 `& (valid > 0)` 的 Triton 语义修正)均为本人独立完成。
94+
- 九齿 0.25.0 的能力边界探针由本人编写和验证,Codex 未参与探针的代码编写。
95+
96+
5. **开源协议类型**:不适用(商业服务)。
97+
98+
6. **其他补充说明**:Codex 作为 AI 编程助手提供代码审查和架构建议,属于参赛规则允许的工具使用范围。
99+
100+
---
101+
102+
## 补充声明
103+
104+
除上述明确标注的参考资源外,本次赛题提交的算子代码的核心算法逻辑及实现方案均为本人在比赛期间独立设计与开发。所有代码均使用九齿 DSL(arrangement/application/premake)或 Triton 语言从零编写,未直接复制任何外部项目的代码片段。

src/ntops/kernels/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
bitwise_not,
88
bitwise_or,
99
bmm,
10+
broadcast,
11+
cartesian_prod,
1012
clamp,
13+
column_stack,
1114
conv2d,
1215
cos,
1316
div,
@@ -24,12 +27,15 @@
2427
lt,
2528
max_pool2d,
2629
mm,
30+
mode,
2731
mul,
2832
ne,
2933
neg,
3034
pow,
35+
repeat,
3136
relu,
3237
rms_norm,
38+
roll,
3339
rotary_position_embedding,
3440
rsqrt,
3541
scaled_dot_product_attention,
@@ -50,6 +56,8 @@
5056
"bitwise_not",
5157
"bitwise_or",
5258
"bmm",
59+
"broadcast",
60+
"cartesian_prod",
5361
"clamp",
5462
"conv2d",
5563
"cos",
@@ -67,10 +75,12 @@
6775
"lt",
6876
"max_pool2d",
6977
"mm",
78+
"mode",
7079
"mul",
7180
"ne",
7281
"neg",
7382
"pow",
83+
"repeat",
7484
"relu",
7585
"rms_norm",
7686
"rotary_position_embedding",

src/ntops/kernels/broadcast.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import functools
2+
3+
from ninetoothed import Symbol, Tensor
4+
5+
6+
def arrangement_dim0(input_1d, output_2d, N):
7+
"""Broadcast input (M,) to (M, N) along dim=0.
8+
9+
output[i, j] = input[i] for all j.
10+
"""
11+
N_sym = Symbol("N", constexpr=True) if not isinstance(N, Symbol) else N
12+
13+
input_arranged = input_1d.tile((-1, 1))
14+
input_arranged = input_arranged.expand((-1, N_sym))
15+
16+
output_arranged = output_2d.tile((-1, -1))
17+
18+
return input_arranged, output_arranged
19+
20+
21+
def arrangement_dim1(input_1d, output_2d, M):
22+
"""Broadcast input (N,) to (M, N) along dim=1.
23+
24+
output[i, j] = input[j] for all i.
25+
"""
26+
M_sym = Symbol("M", constexpr=True) if not isinstance(M, Symbol) else M
27+
28+
input_arranged = input_1d.tile((1, -1))
29+
input_arranged = input_arranged.expand((M_sym, -1))
30+
31+
output_arranged = output_2d.tile((-1, -1))
32+
33+
return input_arranged, output_arranged
34+
35+
36+
def application(input_1d, output_2d):
37+
output_2d = input_1d # noqa: F841
38+
39+
40+
def premake(expand_dim, other_size, dtype=None, block_size=None):
41+
"""Create a broadcast premake function.
42+
43+
Args:
44+
expand_dim: 0 to expand along rows, 1 to expand along columns.
45+
other_size: the size of the other (expanded-to) dimension.
46+
"""
47+
if expand_dim == 0:
48+
arrangement_fn = functools.partial(arrangement_dim0, N=other_size)
49+
elif expand_dim == 1:
50+
arrangement_fn = functools.partial(arrangement_dim1, M=other_size)
51+
else:
52+
raise ValueError(f"expand_dim must be 0 or 1, got {expand_dim}")
53+
54+
tensors = (Tensor(1, dtype=dtype), Tensor(2, dtype=dtype))
55+
56+
return arrangement_fn, application, tensors
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
"""Cartesian product kernel: ntl.load with manual index computation.
2+
3+
Output-driven: for each output row, compute the source index in the input
4+
using integer division and modulo, then ntl.load.
5+
"""
6+
7+
import functools
8+
9+
import ninetoothed
10+
import ninetoothed.language as ntl
11+
from ninetoothed import Tensor
12+
13+
14+
def arrangement(*tensors, block_size=None):
15+
"""Flatten + tile for data tensors (skip 0-dim constexpr)."""
16+
if block_size is None:
17+
block_size = ninetoothed.block_size()
18+
return tuple(
19+
t.flatten().tile((block_size,)) if t.ndim != 0 else t
20+
for t in tensors
21+
)
22+
23+
24+
def application(input, output, repeat_after, size):
25+
"""Compute one column of the cartesian product.
26+
27+
input: 1D tensor, size = L_col
28+
output: 1D column vector, size = total_rows
29+
repeat_after: product of sizes of columns to the right
30+
size: L_col (the size of this input)
31+
32+
source_index = (row // repeat_after) % size
33+
"""
34+
row = output.offsets(0)
35+
src_idx = (row // repeat_after) % size
36+
ptr = input.source.data_ptr() + src_idx * input.source.stride(0)
37+
output = ntl.load(ptr) # noqa: F841
38+
39+
40+
def premake(repeat_after, size, dtype=None, block_size=None):
41+
"""Create cartesian_prod kernel for one input column.
42+
43+
Args:
44+
repeat_after: product(sizes[col+1:]) — how many rows before repeating.
45+
size: L_col — the length of this input tensor.
46+
"""
47+
arrangement_fn = functools.partial(arrangement, block_size=block_size)
48+
tensors = (
49+
Tensor(1, dtype=dtype),
50+
Tensor(1, dtype=dtype),
51+
Tensor(0, constexpr=True, value=repeat_after),
52+
Tensor(0, constexpr=True, value=size),
53+
)
54+
return arrangement_fn, application, tensors

src/ntops/kernels/column_stack.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
"""Column stack kernel: identity copy into sliced output region.
2+
3+
Arrangement: flatten+tile both input and output.
4+
Application: identity — the column offset is handled by output's data pointer
5+
(given by the sliced view from the wrapper).
6+
"""
7+
8+
import functools
9+
10+
import ninetoothed
11+
from ninetoothed import Tensor
12+
13+
14+
def arrangement(*tensors, block_size=None):
15+
"""Flatten + tile for data tensors (skip 0-dim constexpr)."""
16+
if block_size is None:
17+
block_size = ninetoothed.block_size()
18+
return tuple(
19+
t.flatten().tile((block_size,)) if t.ndim != 0 else t
20+
for t in tensors
21+
)
22+
23+
24+
def application(input, output):
25+
"""Identity copy — output already points to correct column range."""
26+
output = input # noqa: F841
27+
28+
29+
def premake(ndim, dtype=None, block_size=None):
30+
"""Create column_stack kernel for one input.
31+
32+
Args:
33+
ndim: ndim of both input and output (2 for column_stack).
34+
"""
35+
arrangement_fn = functools.partial(arrangement, block_size=block_size)
36+
tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype))
37+
return arrangement_fn, application, tensors

0 commit comments

Comments
 (0)