xtl
Back to writing

/ MLC chentianqi

[MLC-02] Tensor Program Abstraction and TensorIR

Notes for lesson two: primitive tensor functions, loops, buffers, block axes, and why TensorIR schedules make tensor programs optimizable.

1 minMLC · TensorIR · Schedule · Tiling

第二讲从“一个算子如何真正执行”出发。计算图知道 matmul 是什么，但若要决定循环如何分块、线程如何映射、数据放在哪一级存储，就需要能暴露程序结构的张量程序抽象。

1. 元张量函数#

模型由许多元张量函数组成，例如：

linear -> add -> relu -> softmax

每个函数可以用库调用、手写循环、向量代码或设备专用实现表达。张量程序抽象关心的是：如何把这些实现表示为可分析、可合法变换的程序。

以矩阵乘法为例：

Cij=∑kAikBkjC_{ij} = \sum_k A_{ik}B_{kj}

最朴素的程序包含三件事：多维缓冲区 A/B/C、i/j/k 循环嵌套，以及归约更新语句。

2. TensorIR 的 block#

TensorIR 不只记录 for-loop，还用 block 标记一个可推理的计算单元：

for i, j, k in T.grid(n, m, r):
    with T.block("C"):
        vi, vj, vk = T.axis.remap("SSR", [i, j, k])
        with T.init():
            C[vi, vj] = 0.0
        C[vi, vj] += A[vi, vk] * B[vk, vj]

SSR 表示两个 spatial axis 和一个 reduce axis。这个额外语义并不是装饰：它告诉编译器哪些循环独立、哪个维度需要归约初始化、读写区域是什么。

结构暴露的信息
Buffershape、dtype、读写位置和存储对象
Loop迭代空间与循环顺序
Block一个逻辑计算单元
Block axisspatial/reduction 语义与索引绑定
Init归约的初始化边界

没有这些信息，编译器很难判断一次 loop reorder 是否仍然正确，或能否把某层循环并行化。

3. Schedule 是等价程序变换#

TensorIR 的核心体验不是“把代码写成另一种语法”，而是对同一计算构造不同实现版本：

block = sch.get_block("C", "main")
i, j, k = sch.get_loops(block)
j0, j1 = sch.split(j, factors=[None, 8])
sch.reorder(i, j0, k, j1)
sch.decompose_reduction(block, k)
sch.vectorize(j1)

这段 schedule 的意义是：先把输出列维度切成块，再调整 loop order，使局部数据复用和向量化成为可能，最后把 reduction init 拆成独立阶段。

常见变换包括：

  • split / fuse：改变迭代空间的分块方式；
  • reorder：改变访问和复用顺序；
  • parallel：将独立 spatial work 映射到并行执行；
  • vectorize：映射到向量化 lane；
  • unroll：展开小循环以降低控制开销；
  • compute_at / reverse_compute_at：重定位中间计算以改善局部性。

4. 为什么它和性能有关#

同一公式的算术量可能相同，但访存次数、缓存命中、线程工作分配完全不同。以矩阵乘法为例，未经分块的程序会反复从更慢存储读取数据；tile 后，一个块可以在局部缓存中复用。

数学等价 != 执行等价
 
性能通常由：访问模式、并行粒度、局部性、向量宽度、同步和寄存器压力共同决定。

TensorIR 的价值是把这些影响性能的结构显式放进 IR，使优化不再只是编译器后端的黑盒猜测。

5. 学习笔记#

这讲让我把“算子优化”理解为：先找出不变量，再只变换实现。block、axis 与读写区域就是这种不变量的载体。后续的 GPU mapping、tensorization 和自动调优，都建立在可被安全变换的张量程序之上。

参考#