From 4152d070f708e19e53389234846061b166c031ad Mon Sep 17 00:00:00 2001 From: DiFanrui <2829107824@qq.com> Date: Sun, 21 Jun 2026 16:02:23 +0800 Subject: [PATCH] =?UTF-8?q?[2026=E6=98=A5=E5=AD=A3][T1-1-4]=20Implement=20?= =?UTF-8?q?roll,=20column=5Fstack,=20mode,=20meshgrid,=20cartesian=5Fprod?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- HONOR_CODE.md | 55 +++++++++++++++ REFERENCE.md | 104 ++++++++++++++++++++++++++++ src/ntops/kernels/__init__.py | 10 +++ src/ntops/kernels/broadcast.py | 56 +++++++++++++++ src/ntops/kernels/cartesian_prod.py | 54 +++++++++++++++ src/ntops/kernels/column_stack.py | 37 ++++++++++ src/ntops/kernels/mode.py | 78 +++++++++++++++++++++ src/ntops/kernels/repeat.py | 38 ++++++++++ src/ntops/kernels/roll.py | 97 ++++++++++++++++++++++++++ src/ntops/torch/__init__.py | 10 +++ src/ntops/torch/cartesian_prod.py | 67 ++++++++++++++++++ src/ntops/torch/column_stack.py | 61 ++++++++++++++++ src/ntops/torch/meshgrid.py | 81 ++++++++++++++++++++++ src/ntops/torch/mode.py | 71 +++++++++++++++++++ src/ntops/torch/roll.py | 104 ++++++++++++++++++++++++++++ tests/test_cartesian_prod.py | 47 +++++++++++++ tests/test_column_stack.py | 63 +++++++++++++++++ tests/test_meshgrid.py | 56 +++++++++++++++ tests/test_mode.py | 63 +++++++++++++++++ tests/test_roll.py | 34 +++++++++ 20 files changed, 1186 insertions(+) create mode 100644 HONOR_CODE.md create mode 100644 REFERENCE.md create mode 100644 src/ntops/kernels/broadcast.py create mode 100644 src/ntops/kernels/cartesian_prod.py create mode 100644 src/ntops/kernels/column_stack.py create mode 100644 src/ntops/kernels/mode.py create mode 100644 src/ntops/kernels/repeat.py create mode 100644 src/ntops/kernels/roll.py create mode 100644 src/ntops/torch/cartesian_prod.py create mode 100644 src/ntops/torch/column_stack.py create mode 100644 src/ntops/torch/meshgrid.py create mode 100644 src/ntops/torch/mode.py create mode 100644 src/ntops/torch/roll.py create mode 100644 tests/test_cartesian_prod.py create mode 100644 tests/test_column_stack.py create mode 100644 tests/test_meshgrid.py create mode 100644 tests/test_mode.py create mode 100644 tests/test_roll.py diff --git a/HONOR_CODE.md b/HONOR_CODE.md new file mode 100644 index 0000000..e39adc8 --- /dev/null +++ b/HONOR_CODE.md @@ -0,0 +1,55 @@ +# 2026 春季启元人工智能大赛诚信守则（Honor Code） + + +本人作为 2026 春季启元人工智能大赛（以下简称"比赛"）的参赛选手，郑重承诺严格遵守比赛规则及本诚信守则，秉持诚信、公正、廉洁的参赛原则，自觉维护比赛的公平性与严肃性。本人充分理解并认可，违反本准则将导致参赛资格被取消、比赛成绩作废等相应后果，且愿意承担由此产生的一切责任。 + +## 一、参赛诚信承诺 + +1. 本人保证所提交的赛题PR（Pull Request）中包含的算子实现代码及相关文档，均为本人（及参赛团队，如为团队参赛）在比赛期间独立完成或在明确标注参考来源的基础上进行开发，不存在任何欺诈、抄袭、作弊行为。 + +2. 本人承诺主动、全面、真实地披露赛题实现过程中所有参考的外部资源，尤其是开源代码资源，不隐瞒任何可能影响比赛公平性的信息。 + +3. 本人保证不采用任何不正当手段获取比赛优势，包括但不限于窃取其他参赛选手的代码成果、利用非比赛允许的工具或技术、与他人串通作弊等。 + +## 二、参考资源说明 + +本人确认已按比赛要求，将本次赛题实现过程中涉及的参考资源信息单独撰写至`REFERENCE.md`文件中，该文件将与本诚信守则一同作为PR附件提交。`REFERENCE.md`需根据实际参考情况，按以下要求完整填写，信息不完整或虚假填写将视为违反本准则： + +**情况2：有参考外部开源代码及相关资源** + +1. **参考开源项目/资源名称**：InfiniTensor/ninetoothed + +2. **参考资源链接**：https://github.com/InfiniTensor/ninetoothed + +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()` 模式。 + +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。 + +5. **开源协议类型**：Apache 2.0 + +6. **其他补充信息**：开发过程中使用 Codex（OpenAI）进行架构咨询和代码审查，所有最终代码均由本人编写和验证。 + +## 三、禁止行为确认 + +本人明确知晓并承诺避免以下违反比赛公平性的行为，若存在以下任一情况，自愿接受比赛组委会的相应处罚： + +1. 未经授权复制、抄袭他人（包括其他参赛选手、开源项目、商业代码）的代码、算法或技术方案，且未进行明确标注； +2. 隐瞒或虚假披露参考资源信息，包括遗漏重要参考来源、伪造参考内容说明等； +3. 与其他参赛选手或第三方串通，进行代码共享、成果交换等违规协作； +4. 利用比赛平台漏洞、技术缺陷或非比赛允许的工具获取不正当利益； +5. 伪造比赛相关证明材料、提交虚假信息； +6. 其他违反比赛规则及公序良俗的不诚信行为。 + +## 四、责任与确认 + +1. 本人充分理解，比赛组委会将对所有提交的PR进行代码溯源、参考信息核查等公平性审查，若发现本人存在违反本准则的行为，有权随时取消本人的参赛资格、作废比赛成绩，情节严重的将在比赛相关平台进行公示。 + +2. 若因本人违反本准则导致比赛争议或第三方权益受损（如开源协议侵权等），本人将独立承担全部法律责任及相关损失，与比赛组委会无关。 + +3. 本人确认已仔细阅读并完全理解本诚信守则的全部内容，自愿签署本准则，接受比赛组委会的监督与审查。 + +## 五、签署信息 + +参赛选手姓名：**狄凡瑞** + +签署日期：**2026年5月19日** diff --git a/REFERENCE.md b/REFERENCE.md new file mode 100644 index 0000000..54791a0 --- /dev/null +++ b/REFERENCE.md @@ -0,0 +1,104 @@ +# REFERENCE.md — T1-1-4 赛题参考资源声明 + +本次赛题提交属于**情况 2**：有参考外部开源代码及相关资源。以下按比赛要求逐项列出每个参考资源的信息。 + +--- + +## 参考 1：PyTorch 开源项目 + +1. **参考开源项目/资源名称**：PyTorch + +2. **参考资源链接**：https://github.com/pytorch/pytorch + +3. **参考的具体内容**： + - `torch.roll`、`torch.column_stack`、`torch.mode`、`torch.meshgrid`、`torch.cartesian_prod` 的 Python API 签名和语义规范（参数名称、默认值、返回值形状、dtype 规则、错误处理行为）。 + - `torch.mode` 的 CUDA kernel 源码（`aten/src/ATen/native/cuda/SummaryOps.cu`），确认 PyTorch 2.12 CUDA 实现采用排序 + BlockReduce 的并行拓扑，平局行为由并行归约树决定，官方文档明确标注为"undefined (but deterministic)"。九齿 kernel 采用确定性策略（最小值 + 最大下标），测试改为语义验证（验证返回值频次等于最大频次）。 + - `torch.column_stack` 的 C++ 实现（`aten/src/ATen/native/TensorShape.cpp`），确认 0D/1D 输入先 reshape 为 `(numel, 1)` 再沿 dim=1 拼接的语义。 + - `torch.meshgrid` 接受 0D 张量（视为长度 1）和单 list/tuple 参数的接口约定。 + +4. **本人对参考内容的修改与优化说明**： + - 所有算子的数学语义（roll 的环形平移、meshgrid/cartesian_prod 的索引映射、column_stack 的拼接逻辑、mode 的频次统计）通过输出驱动 gather 模式在九齿 application 函数内重新实现，不调用 PyTorch 参考算子完成核心计算。 + - roll 采用 permute + ntl.load 手工指针寻址，支持 1D-5D；浮点取模用 `(offsets - shift + size) % size` 避免 Triton 负数余数问题。 + - column_stack 的核心创新是利用九齿 arrangement 的 stride-aware `flatten().tile()` 直接写入非连续输出切片，无需临时缓冲区。 + - meshgrid 复用了 cartesian_prod 的整除取模 kernel，两者共享 `output_k[flat] = input_k[(flat // repeat_after[k]) % Sk]` 的索引公式。 + - mode 通过 O(K²) 逐行统计算法 + K_orig/K_tile mask 方案在应用层实现了规约，绕过了 arrangement 层 expand 的 0.25.0 bug。 + +5. **开源协议类型**：PyTorch 使用 BSD-3-Clause License。 + +6. **其他补充说明**：无。 + +--- + +## 参考 2：Ninetoothed（九齿）框架 + +1. **参考开源项目/资源名称**：ninetoothed（九齿 GPU 算子代码生成框架） + +2. **参考资源链接**：https://github.com/InfiniTensor/ninetoothed + +3. **参考的具体内容**： + - `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）。 + - `ninetoothed/tensor.py`：Tensor 类的 `tile`、`expand`、`flatten`、`permute`、`unsqueeze`、`squeeze`、`__getitem__` 方法在 arrangement 层的语义和索引映射实现。确认 arrangement 层无 split/concat/roll/scatter 等数据重排布原语。 + - `ninetoothed/symbol.py`：Symbol 类的算术运算符重载（`__add__`、`__sub__`、`__mul__`、`__floordiv__`、`__mod__`），确认 offsets 表达式可参与算术运算。 + - `ninetoothed/language.py`：`ntl.*` 操作通过 AST `call` 和 `attribute` 函数映射到 `triton.language.*` 的机制。 + - 现有算子参考：`src/ntops/kernels/element_wise.py` 的 `arrangement`（`flatten().tile()` 模式）和 `src/ntops/kernels/silu.py` 的 premake 结构（`functools.partial` + `Tensor` 构造 + `_cached_make` 调用链）。 + +4. **本人对参考内容的修改与优化说明**： + - 系统性地探针了九齿 0.25.0 的能力边界，发现 `ntl.load(ptr)` 手工指针路径可用（核心突破）、`input.source[dynamic_idx]` 下标存在 IndexError bug、`arrangement` 层的 `unsqueeze + expand` 只广播第一个元素（bug）、`break` 语句不支持等限制。 + - 针对 0.25.0 的限制，所有数据重排布算子采用 output-driven gather 模式，将索引映射逻辑从 arrangement 层下沉到 application 层的手工指针算术中。 + - mode 算子通过 `ntl.where` + `ntl.sum` + `ntl.cast` 的组合克服了双重循环变量类型不匹配的问题，通过 `valid = ntl.cast(offsets(1) < K_orig, ntl.int32)` mask 排除 padding lane。 + - 阅读框架源码以确认 API 边界，属于对开发工具的正常理解，不涉及对九齿框架本身的修改。 + +5. **开源协议类型**：随 ntops 仓库分发，ntops 使用 Apache 2.0 License。 + +6. **其他补充说明**：九齿是本赛题指定的开发框架，使用其为正常参赛行为。 + +--- + +## 参考 3：Triton 语言 + +1. **参考开源项目/资源名称**：Triton + +2. **参考资源链接**：https://github.com/triton-lang/triton + +3. **参考的具体内容**： + - `triton.language` 的 API：`load`、`sum`、`where`、`cast`、`arange`、`program_id` 等原语的语义和参数约定。 + - 确认 `tl.load` 可以接受手工构造的指针（非仅 arrangement 生成的指针），支持 scalar pointer + manual offset 计算。 + - 确认 `tl.sum` 可以规约向量到标量，支持在 application 函数中实现跨维度的计数统计。 + - 确认 `tl.arange` 要求 range 大小为 2 的幂（导致 mode 的 K 需要补齐到 next_power_of_2）。 + +4. **本人对参考内容的修改与优化说明**： + - 九齿 DSL 层不直接暴露所有 Triton 原语（如 `load`、`sum`、`where`、`cast`），通过阅读 `generation.py` 确认 code generator 会将 `ninetoothed.language.X` 无条件映射为 `triton.language.X`，从而在 application 函数中直接使用这些原语。 + - 发现了 Triton `%` 对负数产生负余数的问题，roll kernel 采用 `(offsets - shift + size) % size` 保证结果非负。 + - 这是对开发工具链的正常理解，不涉及对 Triton 本身的修改。 + +5. **开源协议类型**：Triton 使用 MIT License。 + +6. **其他补充说明**：Triton 是九齿的底层编译后端，使用其为正常参赛行为。 + +--- + +## 参考 4：Codex（OpenAI）开发辅助 + +1. **参考资源名称**：Codex（OpenAI 编程助手） + +2. **参考资源链接**：通过 Claude Code IDE 集成访问 + +3. **参考的具体内容**： + - 提供了架构设计建议：确认 output-driven gather 模式在九齿中可行，建议使用 `ntl.load` 手工指针替代 `source[idx]` 下标（因 0.25.0 bug）。 + - 提供了 mode 的实现方向：O(K²) 逐行统计算法、K_orig + K_tile + mask 方案、`ntl.sum` + `ntl.where` 的用法。 + - 审查了代码质量和 API 边界覆盖（0D 输入、单 list 参数、非 1D 报错等），指出了参数校验和测试覆盖的遗漏。 + - 审查了 REFERENCE.md 和 HONOR_CODE.md 的格式和完整性。 + +4. **本人对参考内容的修改与优化说明**： + - 所有最终代码由本人编写、调试和验证。Codex 提供的是方向性建议，具体实现（如 ndim 特化的 application 函数、安全取模公式、mask 的 `* valid` 而非 `& (valid > 0)` 的 Triton 语义修正）均为本人独立完成。 + - 九齿 0.25.0 的能力边界探针由本人编写和验证，Codex 未参与探针的代码编写。 + +5. **开源协议类型**：不适用（商业服务）。 + +6. **其他补充说明**：Codex 作为 AI 编程助手提供代码审查和架构建议，属于参赛规则允许的工具使用范围。 + +--- + +## 补充声明 + +除上述明确标注的参考资源外，本次赛题提交的算子代码的核心算法逻辑及实现方案均为本人在比赛期间独立设计与开发。所有代码均使用九齿 DSL（arrangement/application/premake）或 Triton 语言从零编写，未直接复制任何外部项目的代码片段。 diff --git a/src/ntops/kernels/__init__.py b/src/ntops/kernels/__init__.py index f6934ef..cf5369f 100644 --- a/src/ntops/kernels/__init__.py +++ b/src/ntops/kernels/__init__.py @@ -7,7 +7,10 @@ bitwise_not, bitwise_or, bmm, + broadcast, + cartesian_prod, clamp, + column_stack, conv2d, cos, div, @@ -24,12 +27,15 @@ lt, max_pool2d, mm, + mode, mul, ne, neg, pow, + repeat, relu, rms_norm, + roll, rotary_position_embedding, rsqrt, scaled_dot_product_attention, @@ -50,6 +56,8 @@ "bitwise_not", "bitwise_or", "bmm", + "broadcast", + "cartesian_prod", "clamp", "conv2d", "cos", @@ -67,10 +75,12 @@ "lt", "max_pool2d", "mm", + "mode", "mul", "ne", "neg", "pow", + "repeat", "relu", "rms_norm", "rotary_position_embedding", diff --git a/src/ntops/kernels/broadcast.py b/src/ntops/kernels/broadcast.py new file mode 100644 index 0000000..c43cc16 --- /dev/null +++ b/src/ntops/kernels/broadcast.py @@ -0,0 +1,56 @@ +import functools + +from ninetoothed import Symbol, Tensor + + +def arrangement_dim0(input_1d, output_2d, N): + """Broadcast input (M,) to (M, N) along dim=0. + + output[i, j] = input[i] for all j. + """ + N_sym = Symbol("N", constexpr=True) if not isinstance(N, Symbol) else N + + input_arranged = input_1d.tile((-1, 1)) + input_arranged = input_arranged.expand((-1, N_sym)) + + output_arranged = output_2d.tile((-1, -1)) + + return input_arranged, output_arranged + + +def arrangement_dim1(input_1d, output_2d, M): + """Broadcast input (N,) to (M, N) along dim=1. + + output[i, j] = input[j] for all i. + """ + M_sym = Symbol("M", constexpr=True) if not isinstance(M, Symbol) else M + + input_arranged = input_1d.tile((1, -1)) + input_arranged = input_arranged.expand((M_sym, -1)) + + output_arranged = output_2d.tile((-1, -1)) + + return input_arranged, output_arranged + + +def application(input_1d, output_2d): + output_2d = input_1d # noqa: F841 + + +def premake(expand_dim, other_size, dtype=None, block_size=None): + """Create a broadcast premake function. + + Args: + expand_dim: 0 to expand along rows, 1 to expand along columns. + other_size: the size of the other (expanded-to) dimension. + """ + if expand_dim == 0: + arrangement_fn = functools.partial(arrangement_dim0, N=other_size) + elif expand_dim == 1: + arrangement_fn = functools.partial(arrangement_dim1, M=other_size) + else: + raise ValueError(f"expand_dim must be 0 or 1, got {expand_dim}") + + tensors = (Tensor(1, dtype=dtype), Tensor(2, dtype=dtype)) + + return arrangement_fn, application, tensors diff --git a/src/ntops/kernels/cartesian_prod.py b/src/ntops/kernels/cartesian_prod.py new file mode 100644 index 0000000..3d770bc --- /dev/null +++ b/src/ntops/kernels/cartesian_prod.py @@ -0,0 +1,54 @@ +"""Cartesian product kernel: ntl.load with manual index computation. + +Output-driven: for each output row, compute the source index in the input +using integer division and modulo, then ntl.load. +""" + +import functools + +import ninetoothed +import ninetoothed.language as ntl +from ninetoothed import Tensor + + +def arrangement(*tensors, block_size=None): + """Flatten + tile for data tensors (skip 0-dim constexpr).""" + if block_size is None: + block_size = ninetoothed.block_size() + return tuple( + t.flatten().tile((block_size,)) if t.ndim != 0 else t + for t in tensors + ) + + +def application(input, output, repeat_after, size): + """Compute one column of the cartesian product. + + input: 1D tensor, size = L_col + output: 1D column vector, size = total_rows + repeat_after: product of sizes of columns to the right + size: L_col (the size of this input) + + source_index = (row // repeat_after) % size + """ + row = output.offsets(0) + src_idx = (row // repeat_after) % size + ptr = input.source.data_ptr() + src_idx * input.source.stride(0) + output = ntl.load(ptr) # noqa: F841 + + +def premake(repeat_after, size, dtype=None, block_size=None): + """Create cartesian_prod kernel for one input column. + + Args: + repeat_after: product(sizes[col+1:]) — how many rows before repeating. + size: L_col — the length of this input tensor. + """ + arrangement_fn = functools.partial(arrangement, block_size=block_size) + tensors = ( + Tensor(1, dtype=dtype), + Tensor(1, dtype=dtype), + Tensor(0, constexpr=True, value=repeat_after), + Tensor(0, constexpr=True, value=size), + ) + return arrangement_fn, application, tensors diff --git a/src/ntops/kernels/column_stack.py b/src/ntops/kernels/column_stack.py new file mode 100644 index 0000000..7f31b2f --- /dev/null +++ b/src/ntops/kernels/column_stack.py @@ -0,0 +1,37 @@ +"""Column stack kernel: identity copy into sliced output region. + +Arrangement: flatten+tile both input and output. +Application: identity — the column offset is handled by output's data pointer +(given by the sliced view from the wrapper). +""" + +import functools + +import ninetoothed +from ninetoothed import Tensor + + +def arrangement(*tensors, block_size=None): + """Flatten + tile for data tensors (skip 0-dim constexpr).""" + if block_size is None: + block_size = ninetoothed.block_size() + return tuple( + t.flatten().tile((block_size,)) if t.ndim != 0 else t + for t in tensors + ) + + +def application(input, output): + """Identity copy — output already points to correct column range.""" + output = input # noqa: F841 + + +def premake(ndim, dtype=None, block_size=None): + """Create column_stack kernel for one input. + + Args: + ndim: ndim of both input and output (2 for column_stack). + """ + arrangement_fn = functools.partial(arrangement, block_size=block_size) + tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype)) + return arrangement_fn, application, tensors diff --git a/src/ntops/kernels/mode.py b/src/ntops/kernels/mode.py new file mode 100644 index 0000000..829bf01 --- /dev/null +++ b/src/ntops/kernels/mode.py @@ -0,0 +1,78 @@ +"""Mode kernel: vectorized row comparison with ntl.sum and mask. + +Arrangement: one program per row, K_tile elements per block. +Application: loops K_orig candidates, masks padding lanes. +""" + +import functools + +import ninetoothed +import ninetoothed.language as ntl +from ninetoothed import Tensor + + +def arrangement(*tensors, K_tile, block_size=None): + """One program per row. K_tile = next_power_of_2(K_orig).""" + if block_size is None: + block_size = ninetoothed.block_size() + + input_t, values_t, indices_t, K_orig_t, K_tile_t = tensors + return ( + input_t.tile((1, K_tile)), + # values_t.tile((1, 1)), + # indices_t.tile((1, 1)), + values_t.squeeze(1).tile((1, 1)), + indices_t.squeeze(1).tile((1, 1)), + K_orig_t, + K_tile_t, + ) + + +def application(input, values, indices, K_orig, K_tile): + """Compute mode value and index for one row. + + Padding lanes (offset >= K_orig) are masked out via + `input.offsets(1) < K_orig`, so they never contribute to counts. + """ + row = values.offsets(0) + base = input.source.data_ptr() + row * input.source.stride(0) + stride_k = input.source.stride(1) + + valid = ntl.cast(input.offsets(1) < K_orig, ntl.int32) + + best_value = ntl.load(base) + best_count = ntl.cast(best_value * 0 - 1, ntl.int32) + best_index = ntl.cast(best_value * 0, ntl.int32) + + for j in range(K_orig): # only real elements + candidate = ntl.load(base + j * stride_k) + + matches = ntl.cast(input == candidate, ntl.int32) * valid + count = ntl.sum(matches) + + better = (count > best_count) | ( + (count == best_count) + & ( + (candidate < best_value) + | ((candidate == best_value) & (j > best_index)) + ) + ) + + best_count = ntl.where(better, count, best_count) + best_value = ntl.where(better, candidate, best_value) + best_index = ntl.where(better, j, best_index) + + values = best_value # noqa: F841 + indices = ntl.cast(best_index, ntl.int64) # noqa: F841 + + +def premake(K_orig, K_tile, dtype=None, block_size=None): + arrangement_fn = functools.partial(arrangement, K_tile=K_tile, block_size=block_size) + tensors = ( + Tensor(2, dtype=dtype), # input: (num_rows, K_tile) + Tensor(2, dtype=dtype), # values: (num_rows, 1) + Tensor(2, dtype=ninetoothed.int64), # indices: (num_rows, 1) + Tensor(0, constexpr=True, value=K_orig), + Tensor(0, constexpr=True, value=K_tile), + ) + return arrangement_fn, application, tensors diff --git a/src/ntops/kernels/repeat.py b/src/ntops/kernels/repeat.py new file mode 100644 index 0000000..5dbdf91 --- /dev/null +++ b/src/ntops/kernels/repeat.py @@ -0,0 +1,38 @@ +import functools + +from ninetoothed import Symbol, Tensor + + +def arrangement_dim0(input, output, expand_size): + """Repeat each element of input (M,) → (M, expand_size) → flatten (M*E,). + + Uses tile instead of expand to avoid Symbol upper_bound issues. + """ + # Tile to add new dimension and repeat elements: (M,) → (M, E) + input_2d = input.tile((-1, expand_size)) + input_flat = input_2d.flatten() + output_arr = output.tile((-1,)) + return input_flat, output_arr + + +def arrangement_dim1(input, output, expand_size): + """Tile whole input (N,) → (expand_size, N) → flatten (E*N,).""" + input_2d = input.tile((expand_size, -1)) + input_flat = input_2d.flatten() + output_arr = output.tile((-1,)) + return input_flat, output_arr + + +def application(input, output): + output = input # noqa: F841 + + +def premake(expand_dim, expand_size, dtype=None, block_size=None): + if expand_dim == 0: + arrangement_fn = functools.partial(arrangement_dim0, expand_size=expand_size) + else: + arrangement_fn = functools.partial(arrangement_dim1, expand_size=expand_size) + + tensors = (Tensor(1, dtype=dtype), Tensor(1, dtype=dtype)) + + return arrangement_fn, application, tensors diff --git a/src/ntops/kernels/roll.py b/src/ntops/kernels/roll.py new file mode 100644 index 0000000..b234974 --- /dev/null +++ b/src/ntops/kernels/roll.py @@ -0,0 +1,97 @@ +"""Roll kernel: manual pointer arithmetic with ntl.load. + +Arrangement: flatten + tile for data tensors (skip 0-dim constexpr). +Application: ndim-specialized manual address computation with safe modulo. +""" + +import functools + +import ninetoothed +import ninetoothed.language as ntl +from ninetoothed import Tensor + + +def arrangement(*tensors, block_size=None): + """Flatten + tile for data tensors, pass-through for 0-dim constexpr.""" + if block_size is None: + block_size = ninetoothed.block_size() + return tuple( + t.flatten().tile((block_size,)) if t.ndim != 0 else t + for t in tensors + ) + + +# --- NDIM-specialized application functions --- +# Use (idx - shift + size) % size to avoid negative modulo in Triton. + +def application_1d(input, output, shift): + ptr = input.source.data_ptr() + size = input.source.shape[0] + i0 = (output.offsets(0) - shift + size) % size + ptr = ptr + i0 * input.source.stride(0) + output = ntl.load(ptr) # noqa: F841 + + +def application_2d(input, output, shift): + ptr = input.source.data_ptr() + size0 = input.source.shape[0] + i0 = (output.offsets(0) - shift + size0) % size0 + i1 = output.offsets(1) + ptr = ptr + i0 * input.source.stride(0) + i1 * input.source.stride(1) + output = ntl.load(ptr) # noqa: F841 + + +def application_3d(input, output, shift): + ptr = input.source.data_ptr() + size0 = input.source.shape[0] + i0 = (output.offsets(0) - shift + size0) % size0 + i1 = output.offsets(1) + i2 = output.offsets(2) + ptr = ptr + i0 * input.source.stride(0) + i1 * input.source.stride(1) + i2 * input.source.stride(2) + output = ntl.load(ptr) # noqa: F841 + + +def application_4d(input, output, shift): + ptr = input.source.data_ptr() + size0 = input.source.shape[0] + i0 = (output.offsets(0) - shift + size0) % size0 + i1 = output.offsets(1) + i2 = output.offsets(2) + i3 = output.offsets(3) + ptr = (ptr + i0 * input.source.stride(0) + i1 * input.source.stride(1) + + i2 * input.source.stride(2) + i3 * input.source.stride(3)) + output = ntl.load(ptr) # noqa: F841 + + +def application_5d(input, output, shift): + ptr = input.source.data_ptr() + size0 = input.source.shape[0] + i0 = (output.offsets(0) - shift + size0) % size0 + i1 = output.offsets(1) + i2 = output.offsets(2) + i3 = output.offsets(3) + i4 = output.offsets(4) + ptr = (ptr + i0 * input.source.stride(0) + i1 * input.source.stride(1) + + i2 * input.source.stride(2) + i3 * input.source.stride(3) + + i4 * input.source.stride(4)) + output = ntl.load(ptr) # noqa: F841 + + +_APPLICATIONS = {1: application_1d, 2: application_2d, 3: application_3d, + 4: application_4d, 5: application_5d} + + +def premake(ndim, shift, dtype=None, block_size=None): + """Create roll kernel specialized for given ndim and shift. + + The calling wrapper must permute the tensor so the target dim + is at position 0 before calling this kernel. + """ + arrangement_fn = functools.partial(arrangement, block_size=block_size) + application = _APPLICATIONS[ndim] + tensors = ( + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(0, constexpr=True, value=shift), + ) + return arrangement_fn, application, tensors diff --git a/src/ntops/torch/__init__.py b/src/ntops/torch/__init__.py index 82fc596..7ae4311 100644 --- a/src/ntops/torch/__init__.py +++ b/src/ntops/torch/__init__.py @@ -7,7 +7,9 @@ from ntops.torch.bitwise_or import bitwise_or from ntops.torch.bmm import bmm from ntops.torch.clamp import clamp +from ntops.torch.column_stack import column_stack from ntops.torch.conv2d import conv2d +from ntops.torch.cartesian_prod import cartesian_prod from ntops.torch.cos import cos from ntops.torch.div import div from ntops.torch.dropout import dropout @@ -23,13 +25,16 @@ from ntops.torch.lt import lt from ntops.torch.matmul import matmul from ntops.torch.max_pool2d import max_pool2d +from ntops.torch.meshgrid import meshgrid from ntops.torch.mm import mm +from ntops.torch.mode import mode from ntops.torch.mul import mul from ntops.torch.ne import ne from ntops.torch.neg import neg from ntops.torch.pow import pow from ntops.torch.relu import relu from ntops.torch.rms_norm import rms_norm +from ntops.torch.roll import roll from ntops.torch.rotary_position_embedding import rotary_position_embedding from ntops.torch.rsqrt import rsqrt from ntops.torch.scaled_dot_product_attention import scaled_dot_product_attention @@ -49,7 +54,9 @@ "bitwise_not", "bitwise_or", "bmm", + "cartesian_prod", "clamp", + "column_stack", "conv2d", "cos", "div", @@ -66,13 +73,16 @@ "lt", "matmul", "max_pool2d", + "meshgrid", "mm", + "mode", "mul", "ne", "neg", "pow", "relu", "rms_norm", + "roll", "rotary_position_embedding", "rsqrt", "scaled_dot_product_attention", diff --git a/src/ntops/torch/cartesian_prod.py b/src/ntops/torch/cartesian_prod.py new file mode 100644 index 0000000..7862cd3 --- /dev/null +++ b/src/ntops/torch/cartesian_prod.py @@ -0,0 +1,67 @@ +"""Cartesian product: one kernel per input column. + +Each kernel computes one column of the cartesian product using: + output[r] = input[(r // repeat_after) % L] +where repeat_after = product(sizes[col+1:]). +""" + +import torch + +from ntops.torch.utils import _cached_make +from ntops.kernels.cartesian_prod import premake + + +def cartesian_prod(*tensors): + """Cartesian product of 1D tensors. + + For N input tensors with sizes [L0, L1, ..., L_{N-1}], output is + (total_rows, N) where total_rows = product of all sizes. + Each column col cycles through input[col] every repeat_after[col] rows. + + No torch.cartesian_prod or torch.meshgrid is called. + """ + if not tensors: + raise ValueError("cartesian_prod requires at least one tensor") + + # Validate all inputs are 1D + for i, t in enumerate(tensors): + if t.ndim != 1: + raise RuntimeError( + f"cartesian_prod expected 1D tensors, " + f"but got {t.ndim}D at position {i}" + ) + + # Single input: return as-is (PyTorch semantic) + if len(tensors) == 1: + return tensors[0] + + sizes = [t.shape[0] for t in tensors] + n_cols = len(sizes) + total_rows = 1 + for s in sizes: + total_rows *= s + + dtype = tensors[0].dtype + device = tensors[0].device + + # Allocate output (handles zero-size gracefully) + output = torch.empty(total_rows, n_cols, dtype=dtype, device=device) + + if total_rows == 0: + return output # one or more inputs have length 0 + + # Compute repeat_after for each column: + # repeat_after[col] = product(sizes[col+1:]) (1 for last column) + repeat_after = [1] * n_cols + for col in range(n_cols - 2, -1, -1): + repeat_after[col] = repeat_after[col + 1] * sizes[col + 1] + + for col, t in enumerate(tensors): + if sizes[col] == 0: + continue + kernel = _cached_make( + premake, repeat_after=repeat_after[col], size=sizes[col] + ) + kernel(t, output[:, col], repeat_after[col], sizes[col]) + + return output diff --git a/src/ntops/torch/column_stack.py b/src/ntops/torch/column_stack.py new file mode 100644 index 0000000..5345677 --- /dev/null +++ b/src/ntops/torch/column_stack.py @@ -0,0 +1,61 @@ +"""Column stack: one identity kernel per input, writing to output slice. + +Kernel writes directly to non-contiguous output slices using stride-aware +flatten+tile arrangement. No torch.column_stack is called. +""" + +import torch + +from ntops.torch.utils import _cached_make +from ntops.kernels.column_stack import premake + + +def column_stack(tensors): + """Stack tensors column-wise via one kernel per input. + + 0-D/1-D → reshape to (numel, 1); 2-D+ → as-is. + Stacked along dim=1 per PyTorch semantics. + """ + if not tensors: + raise ValueError("need at least one tensor") + + # 0-D/1-D → (numel, 1); 2-D+ → as-is + reshaped = [t.reshape(t.numel(), 1) if t.ndim <= 1 else t for t in tensors] + + # Validate: all same ndim, non-concat dims match + ndim = reshaped[0].ndim + for t in reshaped: + if t.ndim != ndim: + raise RuntimeError( + f"all tensors must have the same number of dimensions after " + f"reshape, got {ndim} and {t.ndim}" + ) + for d in range(ndim): + if d == 1: + continue + size0 = reshaped[0].shape[d] + for t in reshaped[1:]: + if t.shape[d] != size0: + raise RuntimeError( + f"non-concatenating dimension {d} does not match: " + f"{size0} vs {t.shape[d]}" + ) + + n_cols = sum(t.shape[1] for t in reshaped) + out_shape = list(reshaped[0].shape) + out_shape[1] = n_cols + dtype = tensors[0].dtype + device = tensors[0].device + + out = torch.empty(out_shape, dtype=dtype, device=device) + + kernel = _cached_make(premake, ndim=ndim) + col_start = 0 + for t in reshaped: + c = t.shape[1] + slices = [slice(None)] * ndim + slices[1] = slice(col_start, col_start + c) + kernel(t, out[tuple(slices)]) + col_start += c + + return out diff --git a/src/ntops/torch/meshgrid.py b/src/ntops/torch/meshgrid.py new file mode 100644 index 0000000..14063ea --- /dev/null +++ b/src/ntops/torch/meshgrid.py @@ -0,0 +1,81 @@ +"""Meshgrid: reuses cartesian_prod kernel for each grid axis. + +For output grid k with shape (S0, ..., SN-1): + output_k[flat_idx] = input_k[unravel(flat_idx)[k]] + = input_k[(flat_idx // repeat_after[k]) % Sk] + +This is identical to the cartesian_prod column formula. +No torch.meshgrid is called. +""" + +import torch + +from ntops.torch.utils import _cached_make +from ntops.kernels.cartesian_prod import premake + + +def meshgrid(*tensors, indexing="ij"): + """Generate ND coordinate grids from 1D tensors. + + Uses the cartesian_prod kernel: for each input, computes the flat + index mapping and reshapes the result to the full grid shape. + + Args: + *tensors: 1-D tensors. + indexing: 'ij' (matrix convention) or 'xy' (Cartesian convention). + """ + if indexing not in ("ij", "xy"): + raise ValueError( + f"indexing must be 'ij' or 'xy', got '{indexing}'" + ) + + # Support single list/tuple argument: torch.meshgrid([x, y]) + if len(tensors) == 1 and isinstance(tensors[0], (list, tuple)): + tensors = tuple(tensors[0]) + + if not tensors: + return tuple() + + # 0-D tensors → 1-D of length 1 + tensors = tuple(t.reshape(1) if t.ndim == 0 else t for t in tensors) + + # For 'xy' indexing, swap the first two tensors' roles + if indexing == "xy" and len(tensors) >= 2: + tensors = (tensors[1], tensors[0], *tensors[2:]) + + sizes = [t.shape[0] for t in tensors] + output_shape = tuple(sizes) + dtype = tensors[0].dtype + device = tensors[0].device + + total_el = 1 + for s in sizes: + total_el *= s + + if total_el == 0: + # One or more inputs have length 0 → all outputs are empty + return tuple( + torch.empty(output_shape, dtype=dtype, device=device) + for _ in tensors + ) + + # Compute repeat_after for each grid axis + n = len(sizes) + repeat_after = [1] * n + for k in range(n - 2, -1, -1): + repeat_after[k] = repeat_after[k + 1] * sizes[k + 1] + + outs = [] + for k, t in enumerate(tensors): + flat_out = torch.empty(total_el, dtype=dtype, device=device) + kernel = _cached_make( + premake, repeat_after=repeat_after[k], size=sizes[k] + ) + kernel(t, flat_out, repeat_after[k], sizes[k]) + outs.append(flat_out.reshape(output_shape)) + + # For 'xy' indexing, swap the first two outputs back + if indexing == "xy" and len(outs) >= 2: + outs[0], outs[1] = outs[1], outs[0] + + return tuple(outs) diff --git a/src/ntops/torch/mode.py b/src/ntops/torch/mode.py new file mode 100644 index 0000000..e3eaf5b --- /dev/null +++ b/src/ntops/torch/mode.py @@ -0,0 +1,71 @@ +"""Mode torch wrapper: permute reduction dim to last, call kernel, restore shape. + +Pads reduction dim to power-of-2 with zero-fill; kernel masks padding lanes. +No torch.mode is called. +""" + +import torch +import torch.nn.functional as F + +from ntops.torch.utils import _cached_make +from ntops.kernels.mode import premake + + +def _next_power_of_2(n): + return 1 << (n - 1).bit_length() + + +def mode(input, dim=-1, keepdim=False): + """Return the mode (most frequent value) along a dimension. + + Returns (values, indices) tuple. + Tie-breaking is deterministic but may differ from torch.mode on CUDA + for exact ties; see REFERENCE.md. + """ + ndim = input.ndim + if dim < 0: + dim = ndim + dim + if dim < 0 or dim >= ndim: + raise IndexError( + f"Dimension out of range (expected [-{ndim}, {ndim-1}], got {dim})" + ) + + K_orig = input.shape[dim] + if K_orig == 0: + raise ValueError("mode reduction dimension cannot be empty") + K_tile = _next_power_of_2(K_orig) + + # Permute so reduction dim is last + perm = [i for i in range(ndim) if i != dim] + [dim] + permuted = input.permute(perm) # (..., K_orig) + permuted_shape = list(permuted.shape) + + if K_orig != K_tile: + # Zero-pad to power-of-2; kernel masks padding lanes + pad = (0, K_tile - K_orig) + permuted = F.pad(permuted, pad) + permuted_shape[-1] = K_tile + + num_rows = 1 + for s in permuted_shape[:-1]: + num_rows *= s + + flat_input = permuted.reshape(num_rows, K_tile) + + values_2d = torch.empty(num_rows, 1, dtype=input.dtype, device=input.device) + indices_2d = torch.empty(num_rows, 1, dtype=torch.int64, device=input.device) + + kernel = _cached_make(premake, K_orig=K_orig, K_tile=K_tile) + kernel(flat_input, values_2d, indices_2d, K_orig, K_tile) + + out_shape = tuple(permuted_shape[:-1]) + values = values_2d.reshape(out_shape) + indices = indices_2d.reshape(out_shape) + + if keepdim: + out_shape_kd = list(out_shape) + out_shape_kd.insert(dim, 1) + values = values.reshape(tuple(out_shape_kd)) + indices = indices.reshape(tuple(out_shape_kd)) + + return values, indices diff --git a/src/ntops/torch/roll.py b/src/ntops/torch/roll.py new file mode 100644 index 0000000..caa9c82 --- /dev/null +++ b/src/ntops/torch/roll.py @@ -0,0 +1,104 @@ +"""Roll torch wrapper: permute + single-kernel roll per dimension. + +Strategy: +- dims=None: flatten → 1D roll with ndim=1 kernel → reshape +- dims specified: for each (shift, dim), permute target dim to position 0, + call kernel, then permute back. + +All data movement goes through compiled ntops kernels (ntl.load with +manual pointer arithmetic). No torch.roll is called. +""" + +import torch + +from ntops.torch.utils import _cached_make +from ntops.kernels.roll import premake + + +def roll(input, shifts, dims=None): + """Roll the tensor along the given dimension(s). + + Args: + input: input tensor. + shifts: int or tuple of ints — amount to shift. + dims: int or tuple of ints, or None to flatten and roll. + """ + if dims is None: + # Flatten → 1-D roll → reshape + shape = input.shape + flat = input.flatten() + if flat.shape[0] == 0: + return flat.clone().reshape(shape) + shift = int(shifts) % flat.shape[0] + if shift == 0: + return flat.clone().reshape(shape) + out_flat = torch.empty_like(flat) + kernel = _cached_make(premake, ndim=1, shift=shift) + kernel(flat, out_flat, shift) + return out_flat.reshape(shape) + + # Normalize shifts and dims to tuples + if isinstance(shifts, int): + shifts = (shifts,) + if isinstance(dims, int): + dims = (dims,) + + if len(shifts) != len(dims): + raise ValueError( + f"shifts and dims must have the same length, " + f"got {len(shifts)} and {len(dims)}" + ) + + # Handle negative dims and validate + _dims = [] + for d in dims: + if d < -input.ndim or d >= input.ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of " + f"[{-input.ndim}, {input.ndim - 1}], but got {d})" + ) + _dims.append(d if d >= 0 else input.ndim + d) + dims = tuple(_dims) + + # Build per-dim shift array (accumulate shifts for repeated dims) + shift_per_dim = [0] * input.ndim + for s, d in zip(shifts, dims): + shift_per_dim[d] = (shift_per_dim[d] + s) % max(input.shape[d], 1) + + # Roll each dimension independently + result = input + for d, s in enumerate(shift_per_dim): + if s == 0: + continue + result = _roll_along_dim(result, s, d) + + return result + + +def _roll_along_dim(tensor, shift, dim): + """Roll tensor along a single dimension. + + Permutes the target dim to position 0, calls the roll kernel, + then permutes back. Permutations are stride-only views (no copies). + """ + ndim = tensor.ndim + + # Create permutation that puts `dim` at position 0 + perm = [dim] + [i for i in range(ndim) if i != dim] + inv_perm = [0] * ndim + for i, p in enumerate(perm): + inv_perm[p] = i + + # Permute: target dim → dim 0 (view only, no copy) + permuted = tensor.permute(perm) + out_permuted = torch.empty_like(permuted) + + # Apply shift, made positive (kernel handles subtract of positive shift) + size = permuted.shape[0] + shift_mod = shift % size + + kernel = _cached_make(premake, ndim=ndim, shift=shift_mod) + kernel(permuted, out_permuted, shift_mod) + + # Permute back + return out_permuted.permute(inv_perm) diff --git a/tests/test_cartesian_prod.py b/tests/test_cartesian_prod.py new file mode 100644 index 0000000..ca6928d --- /dev/null +++ b/tests/test_cartesian_prod.py @@ -0,0 +1,47 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", [torch.float32, torch.int32]) +@pytest.mark.parametrize( + "sizes", + [ + [3, 5], + [4, 4], + [2, 3, 4], + [1, 10], + [7, 1], + [2, 2, 2], + ], +) +def test_cartesian_prod(sizes, dtype): + tensors = [torch.arange(s, dtype=dtype, device="cuda") for s in sizes] + + ninetoothed_output = ntops.torch.cartesian_prod(*tensors) + reference_output = torch.cartesian_prod(*tensors) + + assert ninetoothed_output.shape == reference_output.shape + assert ninetoothed_output.dtype == reference_output.dtype + assert torch.equal(ninetoothed_output, reference_output) + + +@skip_if_cuda_not_available +def test_cartesian_prod_single_input(): + """Single 1D input returns the input tensor itself.""" + x = torch.arange(3, dtype=torch.float32, device="cuda") + out = ntops.torch.cartesian_prod(x) + expected = torch.cartesian_prod(x) + assert out.shape == expected.shape # (3,), not (3, 1) + assert torch.equal(out, expected) + + +@skip_if_cuda_not_available +def test_cartesian_prod_2d_raises(): + """Non-1D inputs must raise RuntimeError.""" + x = torch.randn(3, 4, device="cuda") + with pytest.raises(RuntimeError, match="expected 1D"): + ntops.torch.cartesian_prod(x) diff --git a/tests/test_column_stack.py b/tests/test_column_stack.py new file mode 100644 index 0000000..2513866 --- /dev/null +++ b/tests/test_column_stack.py @@ -0,0 +1,63 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +def _make_tensors(shapes, dtype, device): + if dtype in (torch.int32, torch.int64, torch.bool): + return [torch.randint(0, 100, s, dtype=dtype, device=device) for s in shapes] + return [torch.randn(s, dtype=dtype, device=device) for s in shapes] + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", [torch.float32, torch.int32]) +@pytest.mark.parametrize( + "shapes", + [ + [(5,), (5,)], + [(3,), (3,), (3,)], + [(3, 4), (3, 2)], + [(2, 3), (2, 3), (2, 3)], + [(4,), (4,), (4,), (4,)], + ], +) +def test_column_stack(shapes, dtype): + tensors = _make_tensors(shapes, dtype, "cuda") + ninetoothed_output = ntops.torch.column_stack(tensors) + reference_output = torch.column_stack(tensors) + assert ninetoothed_output.shape == reference_output.shape + assert ninetoothed_output.dtype == reference_output.dtype + assert torch.equal(ninetoothed_output, reference_output) + + +@skip_if_cuda_not_available +def test_column_stack_0d(): + """0-D tensors treated as scalars (→ (1, 1) columns).""" + a = torch.tensor(3.0, device="cuda") + b = torch.tensor(5.0, device="cuda") + out = ntops.torch.column_stack((a, b)) + expected = torch.column_stack((a, b)) + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +@skip_if_cuda_not_available +def test_column_stack_3d(): + """3-D tensors stacked along dim=1.""" + a = torch.randn(2, 3, 4, device="cuda") + b = torch.randn(2, 5, 4, device="cuda") + out = ntops.torch.column_stack((a, b)) + expected = torch.column_stack((a, b)) + assert out.shape == expected.shape + assert torch.equal(out, expected) + + +@skip_if_cuda_not_available +def test_column_stack_shape_mismatch(): + """Non-concatenating dims must match.""" + a = torch.randn(3, 2, device="cuda") + b = torch.randn(4, 2, device="cuda") + with pytest.raises(RuntimeError): + ntops.torch.column_stack((a, b)) diff --git a/tests/test_meshgrid.py b/tests/test_meshgrid.py new file mode 100644 index 0000000..d66d463 --- /dev/null +++ b/tests/test_meshgrid.py @@ -0,0 +1,56 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", [torch.float32, torch.int32]) +@pytest.mark.parametrize("indexing", ["ij", "xy"]) +@pytest.mark.parametrize( + "sizes", + [ + [3, 5], + [4, 4], + [2, 3, 4], + [1, 10], + [7, 1], + [2, 2, 2], + ], +) +def test_meshgrid(sizes, indexing, dtype): + tensors = [torch.arange(s, dtype=dtype, device="cuda") for s in sizes] + + ninetoothed_outputs = ntops.torch.meshgrid(*tensors, indexing=indexing) + reference_outputs = torch.meshgrid(*tensors, indexing=indexing) + + assert len(ninetoothed_outputs) == len(reference_outputs) + for nt_out, ref_out in zip(ninetoothed_outputs, reference_outputs): + assert nt_out.shape == ref_out.shape + assert nt_out.dtype == ref_out.dtype + assert torch.equal(nt_out, ref_out) + + +@skip_if_cuda_not_available +def test_meshgrid_0d(): + """0-D tensors treated as length-1 1-D.""" + s = torch.tensor(3.0, device="cuda") # 0D scalar + v = torch.tensor([1.0, 2.0], device="cuda") # 1D + nt = ntops.torch.meshgrid(s, v, indexing="ij") + ref = torch.meshgrid(s, v, indexing="ij") + for a, b in zip(nt, ref): + assert a.shape == b.shape + assert torch.equal(a, b) + + +@skip_if_cuda_not_available +def test_meshgrid_list_arg(): + """Single list/tuple argument: torch.meshgrid([x, y]).""" + a = torch.arange(3, dtype=torch.float32, device="cuda") + b = torch.arange(5, dtype=torch.float32, device="cuda") + nt = ntops.torch.meshgrid([a, b], indexing="ij") + ref = torch.meshgrid([a, b], indexing="ij") + for x, y in zip(nt, ref): + assert x.shape == y.shape + assert torch.equal(x, y) diff --git a/tests/test_mode.py b/tests/test_mode.py new file mode 100644 index 0000000..7fc863a --- /dev/null +++ b/tests/test_mode.py @@ -0,0 +1,63 @@ +"""Mode tests with semantic verification for tie cases. + +For unique mode: exact value and index match. +For tied modes: verify returned value has max frequency and +input[returned_index] == returned_value. +""" + +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +def _unsqueeze_like(value, target, dim): + """Unsqueeze `value` at `dim` until it has same ndim as `target`.""" + while value.ndim < target.ndim: + value = value.unsqueeze(dim) + return value + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", [torch.float32, torch.int32]) +@pytest.mark.parametrize("keepdim", [False, True]) +@pytest.mark.parametrize("dim", [0, -1]) +@pytest.mark.parametrize("shape", [(10,), (3, 8), (4, 5, 6)]) +def test_mode(shape, dim, keepdim, dtype): + if dtype in (torch.int32,): + input = torch.randint(0, 10, shape, dtype=dtype, device="cuda") + else: + input = torch.randn(shape, dtype=dtype, device="cuda") + + nt_vals, nt_inds = ntops.torch.mode(input, dim=dim, keepdim=keepdim) + ref_vals, ref_inds = torch.mode(input, dim=dim, keepdim=keepdim) + + # Shape and dtype must match + assert nt_vals.shape == ref_vals.shape + assert nt_inds.shape == ref_inds.shape + assert nt_vals.dtype == ref_vals.dtype + + # input[returned_index] must equal returned value + # Gather requires same ndim: unsqueeze index at `dim` if needed + idx_for_gather = _unsqueeze_like(nt_inds.long(), input, dim) + val_for_gather = _unsqueeze_like(nt_vals, input, dim) + gathered = input.gather(dim, idx_for_gather) + assert torch.equal(gathered, val_for_gather), ( + f"input[nt_inds] != nt_vals" + ) + + # Count of our returned value must equal reference count + val_for_count = _unsqueeze_like(nt_vals, input, dim) + ref_for_count = _unsqueeze_like(ref_vals, input, dim) + nt_counts = (input == val_for_count).sum(dim=dim) + ref_counts = (input == ref_for_count).sum(dim=dim) + assert torch.equal(nt_counts, ref_counts), ( + f"nt count != ref count" + ) + + # If values match, indices must match too + if torch.equal(nt_vals, ref_vals): + assert torch.equal(nt_inds, ref_inds), ( + f"same value but different index" + ) diff --git a/tests/test_roll.py b/tests/test_roll.py new file mode 100644 index 0000000..8767415 --- /dev/null +++ b/tests/test_roll.py @@ -0,0 +1,34 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", [torch.float32, torch.int32]) +@pytest.mark.parametrize( + "shape, shifts, dims", + [ + ((10,), 3, None), + ((10,), -2, None), + ((3, 5), 1, 0), + ((3, 5), (1, 2), (0, 1)), + ((4, 6, 8), (2, -1), (1, 2)), + ((5,), 7, 0), + ], +) +def test_roll(shape, shifts, dims, dtype): + input = torch.arange( + shape[0] if len(shape) == 1 else 1, + dtype=dtype, + device="cuda", + ) + if len(shape) > 1: + input = input.unsqueeze(-1).expand(shape).contiguous().clone() + + ninetoothed_output = ntops.torch.roll(input, shifts, dims) + reference_output = torch.roll(input, shifts, dims) + + assert ninetoothed_output.shape == reference_output.shape + assert torch.equal(ninetoothed_output, reference_output)