StudyPath Blog

012.2 task2附录：TILELANG_语法入门

20分钟阅读

从 a + b 到 GPU Kernel：TileLang 原理详解

这篇文章面向已经完成 Task2、希望进一步理解实现原理的读者。我们不再按作业步骤填写代码，而是从一维向量加法出发，解释 TileLang 如何描述 GPU 上的数据划分、并行执行和 JIT 编译。

如果你还没有开始编写 solution.py，请先阅读《Task2 TileLang 零基础完成指南》。那份文档会从环境检查开始，带你逐步完成实现、测试和 benchmark。

本文研究的计算仍然很简单：把两个一维向量逐元素相加。

C = A + B

数学并不难。真正需要学习的是：怎样把这件事拆成许多小份，交给 GPU 并行完成。

读完本文后，你应该能够：

  • 说清楚 kernel、tile、block 和尾块分别是什么；
  • 看懂 solution.py 中的 TileLang 代码；
  • 理解为什么既要向上取整，又要做边界判断；
  • 知道测试和 benchmark 分别在验证什么。

本文不会展开 CUDA C++、共享内存、warp 或线程同步；这些概念不是理解当前 kernel 的必要条件。

1. 为什么不直接写 A + B

在 PyTorch 中，向量加法只需要一行：

reference = A + B

这行代码只说明了“算什么”，没有告诉我们 GPU 内部“怎样算”。PyTorch 已经在背后选择并启动了实现加法的 GPU 程序。

Task2 的目的不是取代 PyTorch，而是让我们自己描述这个执行过程：

  1. 把长向量切成若干小段；
  2. 让多个 GPU 工作块同时处理这些小段；
  3. 在每一小段内部并行计算；
  4. 处理最后一段长度不足的情况。

这种运行在 GPU 上、完成某项计算的小程序，通常称为 kernel。

TileLang 提供了一套接近 Python 的领域专用语言（DSL）。我们用 Python 风格的语法描述 kernel，TileLang 再把这份描述编译成 GPU 可以执行的代码。

import tilelang
import tilelang.language as T
  • tilelang 提供 JIT 编译等上层能力；
  • tilelang.language 提供编写 kernel 的语言构件，习惯上简称为 T。

需要特别注意：代码虽然长得像普通 Python，但其中的 T.Kernel、T.Parallel、T.Tensor 等写法具有 GPU 编程语义，不能按普通 Python 循环来理解。

2. 先建立一个分块的心智模型

假设向量长度 N = 10，我们规定每次处理 BLOCK_N = 4 个元素。于是，向量会被切成三段：

全局下标        0  1  2  3 | 4  5  6  7 | 8  9  ×  ×
工作块编号          0       |     1       |     2
块内偏移        0  1  2  3 | 0  1  2  3 | 0  1  2  3

前两段各有 4 个有效元素，最后一段只有 2 个。这个没有填满的最后一段叫作 尾块。

在这份程序中：

  • BLOCK_N 表示每个 tile 计划覆盖多少个元素；
  • pid 表示当前 GPU 工作块的编号；
  • i 表示当前 tile 内部的偏移；
  • pid * BLOCK_N + i 得到向量中的全局下标。

例如，pid = 2、i = 1 时：

index = 2 × 4 + 1 = 9

因此，这次计算处理的是 A[9] + B[9]。

Tile 和 GPU block 是同一个东西吗？

不完全是。

  • Tile 是算法层面的概念：一批一起处理的数据。
  • GPU block 是执行层面的概念：一组协作执行的线程。

在这个简单例子里，一个 GPU block 恰好负责一个 tile，所以两者看起来几乎相同。以后遇到矩阵乘法等复杂 kernel 时，不应再把它们当作同义词。

3. 先读完整代码

下面是 Task2 的核心实现。第一次阅读时，只需要抓住“计算块数—确定下标—边界判断—执行加法”这条主线。

import tilelang
import tilelang.language as T
 
 
@tilelang.jit
def tl_add_1d(A, B, BLOCK_N: int):
    N = T.const("N")
    A: T.Tensor((N,), T.float16)
    B: T.Tensor((N,), T.float16)
    C = T.empty((N,), T.float16)
 
    num_blocks = T.ceildiv(N, BLOCK_N)
 
    with T.Kernel(num_blocks, threads=BLOCK_N) as pid:
        base_idx = pid * BLOCK_N
 
        for i in T.Parallel(BLOCK_N):
            index = base_idx + i
            if index < N:
                C[index] = A[index] + B[index]
 
    return C

把它翻译成自然语言就是：

创建一个长度为 N 的输出向量 C。把输入划分成若干个长度为 BLOCK_N 的 tile，并为每个 tile 启动工作块。每个工作块并行处理自己的元素；只有下标仍在向量范围内时，才读取 A、B 并写入 C。最后返回 C。

接下来逐段拆解。

4. 输入、输出和符号长度

N = T.const("N")
A: T.Tensor((N,), T.float16)
B: T.Tensor((N,), T.float16)
C = T.empty((N,), T.float16)

N 为什么不是一个普通整数？

写 kernel 时，我们希望同一份逻辑能处理长度为 1、127 或 1,048,576 的向量。因此，代码先用：

N = T.const("N")

声明一个名为 N 的符号。真正的长度会在调用时根据输入确定，或在显式编译时通过 N=... 指定。

你可以把符号 N 理解为公式中的未知数：编写公式时先保留名字，使用公式时再代入具体值。

T.Tensor 描述了什么？

A: T.Tensor((N,), T.float16)

这行告诉编译器：

  • A 是一个张量；
  • 形状是 (N,)，也就是长度为 N 的一维向量；
  • 每个元素的类型是 float16。

B 的声明相同。T.empty 则创建一个同样形状和类型的输出缓冲区：

C = T.empty((N,), T.float16)

“empty”表示其中尚未写入有效结果，并不是长度为零。kernel 必须为每个有效位置赋值。

BLOCK_N 是什么？

def tl_add_1d(A, B, BLOCK_N: int):

BLOCK_N 是这份实现的分块参数。例如，BLOCK_N=1024 表示每个 tile 覆盖 1024 个元素。它会影响程序的组织方式和性能，但不应改变计算结果。

5. 为什么块数必须向上取整

num_blocks = T.ceildiv(N, BLOCK_N)

T.ceildiv(a, b) 表示 a / b 后向上取整。

当 N = 10、BLOCK_N = 4 时，需要的块数是：

ceil(10 / 4) = 3

如果误写成普通的向下取整：

10 // 4 = 2

那么下标 8 和 9 将永远没有工作块负责。

这个差异在短向量上更加明显。若 N = 1、BLOCK_N = 1024：

1 // 1024          = 0
ceildiv(1, 1024)   = 1

零个工作块意味着 kernel 什么也不做；一个工作块才能覆盖那个唯一的元素。

可以记住这条规则：

要“装下”全部元素，就向上取整；只取完整的若干组，才向下取整。

6. T.Kernel：每个工作块负责哪一段

with T.Kernel(num_blocks, threads=BLOCK_N) as pid:
    base_idx = pid * BLOCK_N

这里的 with ... as ... 借用了 Python 上下文管理器的外形，但在 TileLang 中，它用来划定 kernel 的执行作用域。

  • num_blocks 指定一维网格中有多少个工作块；
  • threads=BLOCK_N 指定这份练习使用的线程配置；
  • pid 是当前工作块在一维网格中的编号，从 0 开始；
  • base_idx 是当前 tile 的第一个全局下标。

以 BLOCK_N = 4 为例：

pid = 0  →  base_idx = 0  →  负责从下标 0 开始的一段
pid = 1  →  base_idx = 4  →  负责从下标 4 开始的一段
pid = 2  →  base_idx = 8  →  负责从下标 8 开始的一段

接触过 CUDA 的读者可以把 pid 类比为一维的 blockIdx.x。没有 CUDA 基础也不影响理解本文。

T.Kernel 和 T.Parallel 的首字母必须大写。Python 大小写敏感，T.kernel 或 T.parallel 指向的可能是其他对象，不能替代这里的写法。

7. T.Parallel：从块内偏移得到全局下标

for i in T.Parallel(BLOCK_N):
    index = base_idx + i

i 的取值范围是 0 到 BLOCK_N - 1。它表示当前 tile 内部的位置。

虽然语法看起来像 for 循环，但：

for i in range(BLOCK_N):

表示普通 Python 中依次迭代，而：

for i in T.Parallel(BLOCK_N):

是在告诉 TileLang 编译器：这些迭代可以并行执行。编译器会根据目标硬件和线程配置映射这些工作；不要机械地理解成“一个 i 永远对应一个固定的物理线程”。

局部偏移加上当前 tile 的起点，就得到输入向量中的全局位置：

index = pid * BLOCK_N + i

这是整个 kernel 最重要的下标公式。

8. 为什么向上取整后还要判断边界

if index < N:
    C[index] = A[index] + B[index]

向上取整保证了所有有效元素都有人处理，但也让尾块中出现了一些超出范围的下标。

继续使用 N = 10、BLOCK_N = 4 的例子。最后一个工作块会算出：

8, 9, 10, 11

其中只有 8 和 9 有效。若读取 A[10] 或写入 C[11]，就发生了越界访问。其后果可能是结果错误、程序崩溃，甚至错误没有立刻暴露。

因此要先判断：

if index < N:

再执行读取和写入。注意，判断必须包住整个表达式，而不只是 C[index] = ...，因为读取 A[index] 和 B[index] 同样可能越界。

ceildiv 和边界判断解决的是两个互补的问题：

  • ceildiv：不要漏掉尾部的有效元素；
  • index < N：不要碰到尾部之外的无效位置。

二者缺一不可。

9. @tilelang.jit：什么时候发生编译

@tilelang.jit
def tl_add_1d(...):
    ...

@tilelang.jit 是 Python 装饰器。它把函数交给 TileLang 管理，使这段 DSL 描述可以被编译并作为算子调用。

JIT 是 Just-In-Time 的缩写，通常译为“即时编译”。在本任务中有两种使用方式。

直接调用

测试代码这样使用：

out = tl_add_1d(a, b, BLOCK_N=1024)

TileLang 根据输入得到所需信息，编译或复用对应的 kernel，然后在 GPU 上执行并返回 C。

先编译，再反复调用

benchmark 中采用：

kernel = tl_add_1d.compile(N=n, BLOCK_N=1024)
out = kernel(a, b)

这里先为给定的 N 和 BLOCK_N 获得可调用 kernel，再反复执行。这样计时更容易聚焦于 kernel 的运行时间，而不是把首次编译开销混在每次执行中。

可以把整个过程概括为：

TileLang 代码
    ↓ JIT 编译
可执行的 GPU kernel
    ↓ 接收 CUDA Tensor
GPU 并行计算
    ↓
返回输出 Tensor

return C 返回的是输出缓冲区。少了它，直接调用时就无法获得预期结果。

10. 测试用例其实在验证设计

test_add.py 使用了以下长度：

1, 127, 1024, 100003, 1048576

它们分别覆盖了不同的边界情况：

  • 1：输入远小于一个 tile，验证仍会启动工作块；
  • 127：小型尾块，验证边界保护；
  • 1024：恰好填满一个 tile；
  • 100003：较大且不能整除，验证多个完整块之后的尾块；
  • 1048576：较大的规则输入，验证正常规模下的执行。

这体现了测试设计中的一个通用原则：不要只验证“最常见”的输入，还要主动选择能够破坏错误假设的输入。

例如，只测试 N=1024 无法区分下面两种写法：

N // BLOCK_N
T.ceildiv(N, BLOCK_N)

因为二者碰巧都得到 1。加入 N=127 后，错误的向下取整会得到 0，问题立刻暴露。

为什么使用近似比较

测试通过以下代码比较结果：

torch.testing.assert_close(out, a + b, atol=1e-2, rtol=1e-2)

输入类型是 float16。浮点数能够表示的数值有限，运算可能产生舍入误差，因此数值计算通常使用绝对容差和相对容差，而不是要求所有二进制位完全相同。

对于当前的逐元素加法，两条实现路径通常会得到非常接近的结果。更复杂的归约运算会因为加法顺序不同而出现更明显的浮点误差。

11. 如何正确理解 benchmark

benchmark 分别测量 TileLang kernel 与 PyTorch 向量加法的平均耗时。分析结果时，应先区分三类成本：

  1. 编译成本：TileLang 将 DSL 生成可执行 kernel；
  2. 启动成本：CPU 请求 GPU 启动一次 kernel；
  3. 计算与访存成本：GPU 实际读取 A、B，执行加法并写入 C。

benchmark_add.py 先调用 .compile(...)，再对已经编译的 kernel 预热和重复计时，因此主要观察后两类成本。

为什么短向量不一定体现 GPU 的优势

当 N 很小时，真正的加法工作极少，kernel 启动的固定成本可能比计算本身更显著。此时即使 GPU 拥有大量并行执行单元，也没有足够工作来利用它们。

随着 N 增大，访存和计算逐渐占据主要部分。当前向量加法每个元素只做一次加法，却需要读取两个数并写回一个数，因此通常更接近“受内存带宽限制”的算子，而不是“受计算能力限制”的算子。

为什么 BLOCK_N 不是越大越好

增大 BLOCK_N 会改变每个工作块承担的数据量和线程配置。它可能减少工作块数量，也可能受到硬件线程上限、资源占用和调度效率的约束。

因此：

  • 不存在脱离设备和输入规模的最佳 BLOCK_N；
  • 正确做法是先保证结果不随 BLOCK_N 改变，再通过 benchmark 比较；
  • 一次测量可能存在波动，应经过预热并重复多次；
  • PyTorch 已经高度优化，本练习不要求自定义 kernel 一定更快。

性能数据回答的是“在这台设备、这个输入和这套测量方法下发生了什么”，不能直接推广为所有环境下的结论。

12. 这套结构如何迁移到其他算子

向量加法的计算简单，但它已经包含 GPU kernel 的基本骨架：

确定数据划分
      ↓
计算工作块数量
      ↓
每个工作块定位自己的数据
      ↓
块内并行计算
      ↓
保护边界并写回结果

这个骨架可以迁移到更复杂的任务，但每类算子会增加新的问题。

归约

求和或求最大值时，一个输出依赖 tile 中的多个输入。除了并行读取，还需要合并多个中间结果。

Softmax

Softmax 通常先求最大值，再计算指数和总和，最后归一化。它会用到多次遍历、归约和数值稳定性处理，但仍然需要定位 tile 和保护边界。

矩阵乘法

矩阵乘法会把二维矩阵切成小块，并让多个工作块处理不同的输出区域。此时 tile 与 GPU block 的区别更加重要，还会涉及数据复用和更复杂的内存层次。

Task2 的价值正在于此：先在没有复杂数学干扰的情况下，掌握数据划分、并行下标和边界处理。

13. 总结

这份 kernel 的核心可以归纳为四点：

  1. 分块：把长度为 N 的输入切成大小为 BLOCK_N 的 tile；
  2. 定位：用 pid * BLOCK_N + i 得到全局下标；
  3. 覆盖：用 T.ceildiv 确保尾部元素不会被漏掉；
  4. 保护：用 index < N 阻止尾块越界。

在语法层面，T.Kernel 描述工作块网格，T.Parallel 表达可以并行的迭代，@tilelang.jit 则把这份描述连接到编译和执行过程。

理解这些内容后，可以继续阅读 TileLang Language Basics，并尝试回答三个问题：

  1. 如果去掉 T.ceildiv，哪些测试会失败，为什么？
  2. 如果保留 T.ceildiv 但去掉边界判断，会访问哪些非法下标？
  3. 为什么改变 BLOCK_N 可能影响速度，却不应该改变结果？

关系图谱

反向链接