StudyPath Blog

012.4 task2附录：NINETOOTHED_语法入门

12分钟阅读

从 Tile 到 Kernel：NineToothed 原理详解

本文面向已经完成 Task2 NineToothed 向量加法的读者。目标是让你理解为什么 NineToothed 仅靠“切分三个向量”和“一行加法”就能形成 GPU kernel。

如果你还没有完成 ninetoothed_add.py，请先阅读《Task2 NineToothed 零基础完成指南》。

1. NineToothed 的核心思路：安排，再计算

对于同一个向量加法：

output = lhs + rhs

TileLang 让你直接描述工作块、并行循环和全局下标；NineToothed 则把问题拆成两个层次：

arrangement：数据如何切块、不同参数的 tile 如何对应
application：一组对应 tile 具体做什么计算

这种模式称为 arrange-and-apply。

在 Task2 中：

def arrangement(lhs, rhs, output):
    return (
        lhs.tile((BLOCK_SIZE,)),
        rhs.tile((BLOCK_SIZE,)),
        output.tile((BLOCK_SIZE,)),
    )
 
 
def application(lhs, rhs, output):
    output = lhs + rhs

arrangement() 决定“哪些数据一起交给一个程序”；application() 决定“这个程序拿到数据后如何计算”。

2. 三种“张量”视角

初学 NineToothed 时，最容易困惑的是：同一个名字 lhs，到底代表完整的 PyTorch 张量，还是一个 tile？

答案取决于代码所处的阶段。

调用时：真实的 PyTorch Tensor

在 nt_add_1d() 中：

def nt_add_1d(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor:
    output = torch.empty_like(lhs)
    _KERNEL(lhs, rhs, output)
    return output

lhs、rhs 和 output 是含有真实数值、位于 GPU 上的 PyTorch Tensor。

构造 kernel 时：Tensor(1) 是符号描述

_KERNEL = ninetoothed.make(
    arrangement,
    application,
    (Tensor(1), Tensor(1), Tensor(1)),
)

Tensor(1) 不保存实际向量元素。它表示“这里将传入一个一维张量”，让 NineToothed 能在不知道具体长度的情况下分析 arrangement() 和 application()。

这种不包含真实数据、只描述维度和布局的对象，可以称为符号张量。

application() 中：当前程序接收到的是 tile

经过 .tile((BLOCK_SIZE,)) 后，完整向量被组织为许多个 tile。application(lhs, rhs, output) 的三个参数不再是完整向量，而是位置相同的一组 tile。

因此：

output = lhs + rhs

不是对整个输入执行一次普通 PyTorch 加法，而是描述“每个程序都对自己拿到的 tile 执行逐元素加法”。

3. .tile() 如何把一维向量变成分层张量

考虑长度 N = 10 的一维向量，令 BLOCK_SIZE = 4：

原始向量：
[ 0  1  2  3  4  5  6  7  8  9 ]

执行：

x_tiled = x.tile((4,))

可以把结果想成两层：

外层：有 3 个 tile
内层：每个 tile 的计划形状是 4
 
x_tiled[0] → 原始下标 0、1、2、3
x_tiled[1] → 原始下标 4、5、6、7
x_tiled[2] → 原始下标 8、9 和尾部位置

外层的长度是向上取整后的 tile 数：

ceil(10 / 4) = 3

内层的形状是 (4,)。这就是“分层张量”：外层元素本身是一个 tile，而不是单个标量。

尾块为什么不需要手写边界判断

Tensor.tile() 的默认 floor_mode=False，会保留不能填满的尾块，而不是直接丢弃它。因此，N=10、BLOCK_SIZE=4 仍会产生第 3 个 tile。

NineToothed 根据 tile 布局生成 kernel 时处理尾块的有效边界。Task2 的实现只要让三个向量以同样方式调用 .tile((BLOCK_SIZE,))，就可以覆盖尾部元素。

这与 TileLang 的实现方式不同：TileLang 要显式计算下标并写 if index < N；NineToothed 把这部分边界处理封装在 tile 布局和生成过程里。两者要解决的问题相同，抽象层次不同。

4. arrangement() 真正做了什么

def arrangement(lhs, rhs, output):
    return (
        lhs.tile((BLOCK_SIZE,)),
        rhs.tile((BLOCK_SIZE,)),
        output.tile((BLOCK_SIZE,)),
    )

这段代码没有执行加法，也没有复制真实数据。它在构造 kernel 时说明三个参数的布局关系。

因为三个向量长度相同、tile 大小相同，它们的外层形状也相同：

lhs 外层：tile 0、tile 1、tile 2
rhs 外层：tile 0、tile 1、tile 2
output 外层：tile 0、tile 1、tile 2

NineToothed 会据此为每个外层位置启动一个程序，并把相同位置的 tile 配对：

程序 0：lhs tile 0、rhs tile 0 → output tile 0
程序 1：lhs tile 1、rhs tile 1 → output tile 1
程序 2：lhs tile 2、rhs tile 2 → output tile 2

这解释了为什么三个参数必须：

  1. 都进行 tile；
  2. 使用兼容的 tile 布局；
  3. 按 lhs、rhs、output 的函数参数顺序返回。

对于向量加法，三者的布局完全相同就足够了。矩阵乘法等任务会需要 tile、expand、squeeze 等额外操作，使不同形状的输入 tile 也能正确对应到同一个输出 tile。

5. application() 为什么只有一行

def application(lhs, rhs, output):
    output = lhs + rhs

在这一层，lhs、rhs 和 output 都是当前程序取得的 tile。它们的形状相同，例如都是 (1024,)。

因此，lhs + rhs 表示当前 tile 内对应位置相加：

lhs tile = [a0, a1, a2, ...]
rhs tile = [b0, b1, b2, ...]
结果     = [a0+b0, a1+b1, a2+b2, ...]

NineToothed 的 application 语义与普通 Python 函数不同。此处：

output = lhs + rhs

是对编译器的描述：把当前 tile 的结果写入对应输出 tile。它不需要 return output，也不应把它改成只改变 Python 局部变量的普通控制流程。

也不需要手写：

for i in range(BLOCK_SIZE):
    output[i] = lhs[i] + rhs[i]

因为应用阶段已经是 tile 级计算，NineToothed 会把逐元素操作编译成底层 GPU 操作。

6. ninetoothed.make() 如何连接两阶段

_KERNEL = ninetoothed.make(
    arrangement,
    application,
    (Tensor(1), Tensor(1), Tensor(1)),
)

可以将它理解为一个组装动作：

一维参数描述
    + arrangement 的布局规则
    + application 的 tile 计算规则
    ↓
可调用的 GPU kernel：_KERNEL

调用 _KERNEL(lhs, rhs, output) 时，真实 PyTorch Tensor 会按照 arrangement() 描述切分，随后对每一组对应 tile 执行 application() 描述的操作。

nt_add_1d() 的价值是提供普通 PyTorch 风格的接口：

out = nt_add_1d(lhs, rhs)

调用者不需要自己分配 output 或直接管理 _KERNEL。

7. BLOCK_SIZE 与自动调优

文件默认写法为：

AUTOTUNE = os.environ.get("NINETOOTHED_AUTOTUNE") == "1"
BLOCK_SIZE = (
    ninetoothed.block_size(lower_bound=256, upper_bound=1024)
    if AUTOTUNE
    else 1024
)

默认情况下，BLOCK_SIZE 是固定的 1024。开启：

NINETOOTHED_AUTOTUNE=1 python benchmark_ninetoothed_add.py

后，ninetoothed.block_size(...) 表示让编译器在给定范围内搜索配置。

这不是在运行时按每个元素修改 tile 大小，而是在构建和测量不同 kernel 配置后选择合适的 BLOCK_SIZE。环境变量必须在导入 ninetoothed_add.py 前设置，因为模块导入时就会确定 BLOCK_SIZE 的定义方式。

为什么不能先认定更大的 block 更快

更大的 tile 会减少 tile 数量，但也会改变每个程序承担的工作量和资源使用方式。更小的 tile 则提供更多程序，有时更容易利用硬件并行度。

没有对所有 GPU 和所有输入长度都最优的固定值。benchmark 的意义是观察当前环境中的实际结果，而不是证明某一个数字永远最好。

自动调优与 BLOCK_SIZE 的实测对比记录见 调参实验。

8. 测试用例关注的边界

测试使用：

size = 98432
BLOCK_SIZE = 1024

它们的关系是：

98432 = 96 × 1024 + 128

因此会产生 96 个完整 tile 和一个只含 128 个有效元素的尾块。

如果错误地让 tile 只保留完整块，最后 128 个元素将没有被计算。这个测试验证的正是“相同布局 + 尾块处理”的正确性，而不只是普通长度的向量加法。

9. 与 TileLang 的差异应该怎样理解

两种实现都需要解决同样的问题：覆盖全部元素、把对应输入对齐、并行执行计算、正确处理尾块。

区别主要在谁负责表达这些细节：

TileLang：
你显式写工作块数量、全局下标、并行循环和边界判断。
 
NineToothed：
你声明张量的 tile 布局和 tile 级操作，框架据此构建映射与边界处理。

这不意味着 NineToothed “不使用并行”或 “没有 GPU block”。它表示这些底层映射被放在框架的安排与生成阶段，而不是由本题代码直接写出。

10. 从向量加法迁移到后续算子

Task2 的向量加法只有一层 tile，且每个输出元素只依赖相同位置的两个输入元素，所以 application() 可以是一行。

后续任务会逐步增加复杂性：

  • 逐元素算子：保持相同布局，改变 tile 内表达式；
  • 归约：一个输出值依赖 tile 中多个元素，需要 sum、max 等操作；
  • 矩阵乘法：每个输出 tile 依赖多个输入 tile，需要在 arrangement 中对齐并在 application 中迭代；
  • 融合算子：在一次 application 中完成多个连续操作，减少中间结果的读写。

无论复杂度如何，先问两个问题通常最有帮助：

  1. 每个输出 tile 需要哪些输入 tile？
  2. 怎样在 arrangement() 中让它们对应到同一个程序？

11. 总结

Task2 的 NineToothed 实现很短，是因为工作被清晰地分给了两个阶段：

  1. .tile((BLOCK_SIZE,)) 将向量组织为对应的 tile；
  2. output = lhs + rhs 指定每组 tile 的计算；
  3. ninetoothed.make(...) 将布局与计算组装成可调用 kernel；
  4. .tile() 的默认分块语义保留尾块，保证不能整除时不会漏算。

建议接着阅读官方的 The Basics 和 Tensor.tile API。阅读时可以检验自己是否能回答：

  1. 为什么 application() 的参数不是完整的原始向量？
  2. 为什么三个向量必须具有兼容的外层 tile 形状？
  3. 为什么本题不需要手写 index < N？
  4. 自动调优实际在选择什么？

关系图谱

反向链接