/ MLC chentianqi

[MLC-06] GPU 硬件加速与张量化

陈天奇 MLC 课程第六讲学习笔记：从 thread block、local/shared memory tiling 到硬件专业化和 tensorization，理解张量程序如何映射到 GPU。

2026年7月13日2 minMLC · GPU · CUDA · Tensorization

第六讲把 TensorIR schedule 落到 GPU 执行模型上，并进一步讨论为什么新硬件需要 blockization、特殊 memory scope 与 tensorization。

1. 从循环到线程块#

一个向量加法的逻辑循环：

for i in range(1024):
    C[i] = A[i] + B[i]

可以先 split，再绑定到 GPU 的 block 和 thread：

i0, i1 = sch.split(i, factors=[None, 128])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")

blockIdx.x  选择一个 thread block
threadIdx.x 选择 block 内的一个线程

这并不是语法替换。它决定了并行粒度、线程数、索引计算和后续 memory coalescing 的可能性。

朴素矩阵乘法：

C_{ij} = \sum_k A_{ik}B_{kj}

会频繁从全局内存读取 $A$ 、 $B$ 。GPU 计算单元通常比外部内存快得多，因此优化的第一原则是：让被重复使用的数据停留在更近的存储层级。

课程逐步引入两类分块：

方法	数据位置	主要收益
Local blocking	单线程局部寄存器/局部缓存	减少单线程重复读取
Shared-memory blocking	同一 block 共享存储	多线程复用同一个 tile

一个典型 tile 路径是：

global memory
  -> shared tile of A/B
  -> per-thread register fragment
  -> accumulate C fragment
  -> global memory

TensorIR 的 cache 和 compute placement 可以显式表达数据移动：

A_shared = sch.cache_read(block_C, 0, "shared")
B_shared = sch.cache_read(block_C, 1, "shared")
C_local = sch.cache_write(block_C, 0, "local")

随后再通过 compute_at、loop split、binding 与 cooperative fetch 把复制工作放到合适的线程和循环层级。

正确性之外还要考虑：

现代硬件不只执行标量 FMA，还可能提供 tile 级矩阵指令和专门的搬运指令。此时编译器需要把一段循环计算识别为一个更大的计算单元：

scalar loop nest
    -> blockize a tile computation
    -> map the block to a tensor intrinsic
    -> map buffers to special memory scopes

TensorIR block 是完成这件事的重要桥梁。它既能表示元素级计算，也能表示一个 16 x 16 x 16 的矩阵块；后者可以被替换为硬件或库提供的专用实现。

tensorization != "自动使用更快指令"
 
它要求：计算形状、数据布局、内存层级、循环边界和 intrinsic 契约全部匹配。

GPU 优化的本质不是把更多线程塞进 kernel，而是把数据复用、并行工作和硬件计算单元匹配起来。课程从 threadIdx.x 到 shared tile，再到 tensorized block 的顺序很合理：每一步都在增加 IR 中可被后端利用的结构信息。