StudyPath Blog

012.3 task2附录：NINETOOTHED_零基础完成指南

9分钟阅读

Task2 NineToothed 零基础完成指南

这份指南为你提供完成 task2 NineToothed 任务 需要的前置知识。

NineToothed 的写法和 TileLang 不同。它把一个 kernel 分成两步：

  1. arrangement()：说明输入、输出如何切成 tile；
  2. application()：说明一组对应 tile 如何计算。

先跟着本文完成任务；完成后再阅读《从 Tile 到 Kernel：NineToothed 原理详解》，理解符号张量、分层张量和自动调优。

0. 任务到底要完成什么

输入是两个同长度的一维 float16 向量，输出是逐元素相加的结果：

lhs = [1, 2, 3]
rhs = [4, 5, 6]
output = [5, 7, 9]

也就是：

output[i] = lhs[i] + rhs[i]

你需要：

  1. 完成 ninetoothed_add.py 的 arrangement()；
  2. 完成 ninetoothed_add.py 的 application()；
  3. 处理长度不能被 BLOCK_SIZE 整除的尾块；
  4. 通过 test_ninetoothed_add.py；
  5. 运行 benchmark_ninetoothed_add.py 并保存结果；
  6. 用 3～5 句话记录分块方式和开发体验。

1. 先确认运行环境

以下命令需要在配置好的 GPU 环境中运行。进入任务目录并安装 NineToothed：

cd /data/gollamago/assignment/task2
python -m pip install ninetoothed

然后确认 Python 能找到 NineToothed，且 PyTorch 能看到 GPU：

python -c "import ninetoothed, torch; print('NineToothed: OK'); print('GPU:', torch.cuda.is_available())"

预期包含：

NineToothed: OK
GPU: True

若出现 No module named ninetoothed，请确认安装命令使用的是当前运行测试的同一个 Python：使用 python -m pip，不要只写 pip。若 GPU: False，先检查课程环境，不要先修改 kernel。

2. 认识要修改的文件

打开 ninetoothed_add.py。你只需要补全两个函数：

def arrangement(lhs, rhs, output):
    # TODO(student): tile lhs, rhs, and output with BLOCK_SIZE.
    raise NotImplementedError("Complete arrangement().")
 
 
def application(lhs, rhs, output):
    # TODO(student): compute the tile-wise vector addition.
    raise NotImplementedError("Complete application().")

其余代码已经把你的两个函数组装成可调用的 kernel：

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

现阶段只需知道：

  • lhs、rhs 是两个输入向量；
  • output 是待写入的输出向量；
  • Tensor(1) 表示这三个参数都是一维张量的描述，不是实际数据；
  • _KERNEL 会根据 arrangement 和 application 生成 kernel；
  • nt_add_1d() 已经负责创建真实的 PyTorch 输出张量并调用 _KERNEL。

不要修改 nt_add_1d() 或 _KERNEL 的参数顺序。它们与两个函数的参数顺序必须一致：lhs、rhs、output。

3. 先建立“切 tile，再计算”的心智模型

假设：

向量长度 N = 10
BLOCK_SIZE = 4

把一个向量切成长度为 4 的 tile 后，可以按下面理解：

第 0 个 tile：下标 0、1、2、3
第 1 个 tile：下标 4、5、6、7
第 2 个 tile：下标 8、9，以及未填满的剩余位置

lhs、rhs 和 output 都用相同方式切分。因此每次计算接收到的是一组对应 tile：

第 0 次：lhs[0:4]、rhs[0:4]  → output[0:4]
第 1 次：lhs[4:8]、rhs[4:8]  → output[4:8]
第 2 次：lhs[8: ]、rhs[8: ]  → output[8: ]

最后一个 tile 没有填满，叫作尾块。

本题中不需要像 TileLang 那样手写 index < N。.tile() 默认保留尾块，NineToothed 在生成 kernel 时会处理最后一个不完整 tile 的有效边界。你的职责是让三个向量以相同的 BLOCK_SIZE 切分。

4. 第一步：在 arrangement() 中切分三个向量

将 arrangement() 函数体替换为：

return (
    lhs.tile((BLOCK_SIZE,)),
    rhs.tile((BLOCK_SIZE,)),
    output.tile((BLOCK_SIZE,)),
)

.tile((BLOCK_SIZE,)) 做了什么

它表示把一维向量切成长度为 BLOCK_SIZE 的连续小块。

lhs.tile((BLOCK_SIZE,))

括号中的逗号很重要：(BLOCK_SIZE,) 是 Python 中“只含一个元素的元组”，这里用它表示 tile 的形状是一维、长度为 BLOCK_SIZE。

为什么三个参数都要切，并且顺序不能变

arrangement() 的返回值分别对应传入的三个参数：

第 1 项 → lhs
第 2 项 → rhs
第 3 项 → output

当三个向量都按同样大小切分时，NineToothed 能把第 k 个 lhs tile、第 k 个 rhs tile 和第 k 个 output tile 交给同一次计算。

不要只切输入而遗漏 output，也不要把返回顺序改成 rhs, lhs, output。后者虽然加法本身满足交换律，但会破坏“函数参数和返回 tile 一一对应”的清晰约定。

BLOCK_SIZE 从哪里来

文件已经定义好了它：

BLOCK_SIZE = (
    ninetoothed.block_size(lower_bound=256, upper_bound=1024)
    if AUTOTUNE
    else 1024
)

默认值是 1024。当用环境变量开启自动调优时，NineToothed 会在 256 到 1024 的范围内选择配置。完成基本功能前不需要修改这里的代码。

5. 第二步：在 application() 中计算一组 tile

将 application() 函数体替换为：

output = lhs + rhs

这行代码完成的是一组 tile 的逐元素加法，而不是完整原始向量的加法。比如当前处理第 1 个 tile 时，它等价于：

output[4:8] = lhs[4:8] + rhs[4:8]

所有 tile 都按同样规则处理后，整个向量加法就完成了。

这里有两个容易犯错的地方：

  • 不要再写 Python 的 for 循环。 tile 内部的逐元素计算由 NineToothed 生成的 kernel 并行完成；
  • 不要写 return output。 application() 中的 output = lhs + rhs 是 NineToothed DSL 对“把结果写入当前输出 tile”的描述；它看起来像 Python 赋值，但不需要按普通函数那样返回。

6. 完成后的核心代码

两个函数合起来应当是：

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

确认已经删除两个 raise NotImplementedError(...)。只要其中任意一个仍会执行，测试就会失败。

7. 运行测试并理解结果

在任务目录执行：

python -m pytest -q test_ninetoothed_add.py

预期结果类似：

1 passed

测试长度为 98432，而默认 BLOCK_SIZE=1024：

98432 = 96 × 1024 + 128

因此测试特意包含一个只有 128 个有效元素的尾块。通过测试意味着你的实现不仅能处理完整 tile，也能正确处理最后的不完整 tile。

常见报错

仍然出现 NotImplementedError

检查两个函数中的占位异常是否都已删除。

结果不正确或 kernel 生成失败

优先检查：

  1. 三个 .tile((BLOCK_SIZE,)) 是否都存在；
  2. arrangement() 是否按 lhs、rhs、output 顺序返回；
  3. application() 是否准确写为 output = lhs + rhs。

No module named ninetoothed

在当前终端再次运行：

python -m pip install ninetoothed

8. 运行 benchmark

测试通过后执行：

python benchmark_ninetoothed_add.py

它会对 2^18 到 2^27 的向量长度逐一验证正确性，并输出 NineToothed 和 PyTorch 的耗时。确认每一行的 correct 列均为 PASS，再保存终端截图。

如果要尝试自动调优，在启动 Python 前设置环境变量：

NINETOOTHED_AUTOTUNE=1 python benchmark_ninetoothed_add.py

这个变量必须和命令写在同一行，或在启动脚本之前设置；因为 BLOCK_SIZE 在导入 ninetoothed_add.py 时就已经决定。

不要只根据一次运行的速度判断优劣。不同输入长度和设备可能得到不同结果；本任务首先要求结果正确。

9. 完成检查清单

  • arrangement() 返回了三个按 BLOCK_SIZE 切分的 tile；
  • 三个 tile 的返回顺序是 lhs、rhs、output；
  • application() 使用 output = lhs + rhs；
  • 两个 NotImplementedError 都已删除；
  • python -m pytest -q test_ninetoothed_add.py 通过；
  • python benchmark_ninetoothed_add.py 中每一行均为 PASS；
  • 已保存 benchmark 终端截图；
  • 已用 3～5 句话记录分块方式和实际开发体验。

至此，Task2 的 NineToothed 部分已完成。接下来阅读《从 Tile 到 Kernel：NineToothed 原理详解》，再理解为何 .tile() 能形成分层张量、ninetoothed.make() 如何生成 kernel，以及自动调优实际调整了什么。

关系图谱

反向链接