想象你是一家快递分拣中心的调度员:仓库在园区最外侧,开车往返一趟要十五分钟;流水线工作台旁只有一个放不下多少货的中转货架;而每位打包员手边只有一块贴身的小托盘。一次性把订单全搬过去,货架瞬间爆仓;一趟只带一件,油费和人工又白白浪费。聪明的调度员会把大订单切成大小合适的小批次:先算好每车装多少、按什么顺序装、卸到货架的哪个格子、再分给哪几位打包员——这套“分批搬运”的学问,正是 TVM 中的算子切分与 Tiling 分块策略。它要解决的核心问题只有一个:让数据尽量靠近计算,让每一次远途运输都物超所值。
掌握分块策略,是 AI 编译器工程师优化算子性能的看家本领:它是打通“并行度、访存与计算”三者平衡的关键,直接决定模型在 GPU 上的推理速度能否再上一个台阶。
可视化演示
动态演示TVM算子切分与Tiling的核心机制,通过基础/核心/应用/扩展四种模式展示关键概念与交互流程。
本节配有交互动效(进入视口后自动加载,约 3KB)。
TVM算子切分与Tiling:算子切分 算子切分与Tiling核心概念 GPU缓存层次结构与访存优化 核心机制可视化 TVM算子切分与Tiling整体架构与关键组件
访存层次基础
假设你是一家快递分拣中心的调度员。仓库在园区最外侧,开车往返要十几分钟;工作台旁的中转货架只有几步远;而打包员贴身的托盘一抬手就能摸到。任何调度方案都必须尊重这个物理现实:越近的地方越贵、越小,越远的地方越便宜、越大。你每天的工作,就是在这三个距离之间寻找最优解。
GPU 的存储布局几乎复刻了这座分拣中心。全局内存(global memory)容量以 GB 计,但一次访存动辄数百个时钟周期;共享内存(shared memory)是线程块内所有线程共用的“货架”,时延缩短到几十个周期;寄存器则是每个线程私有的“托盘”,几乎零延迟。TVM 做算子优化的第一步,就是为数据的搬运路线——从全局到共享、再到寄存器——安排好位置,这些位置在 TVM 中被称为存储作用域(storage scope)。
判断数据该放在哪一层,不能凭感觉,要看“有效访存时延”:它把远路和近路的成本按比例加权,算出一笔总账。
调度员心里始终装着一张“距离表”:远路要尽量少走,近路要尽量装满。后续的 Tiling 分块、共享内存缓存、向量化,本质上都是围绕这张距离表在精打细算——把大部分访存从“远路”挪到“近路”。
- 有效时延$t_{\text{有效}}$ 表示一次逻辑访存平均消耗的周期数,相当于调度员眼中“取一次货的平均时间”。
- 远路比例$p$ 是穿透到全局内存的访问占比,等于“远路次数 ÷ 总访问次数”,越接近 0 说明数据越贴近计算。
- 快慢层代价$t_{\text{全局}}$ 与 $t_{\text{共享}}$ 分别是远路和近路的单次耗时,通常相差近 20 倍。
- 优化方向让 $p$ 趋于 0,整体时延就趋近于近路成本。
切分设计动机
为什么不能一次性把订单全搬出来?因为中转货架就那么大,硬塞只会爆仓;为什么不能一件一件搬?因为运输通道太贵,来回跑五百趟谁都受不了。聪明的调度员会找一个中间值:每车装多少件,既放得下货架,又值得跑一趟——这就是“切分”最初的设计动机。
未分块的朴素算子循环正是“一件一件搬”的典型。以矩阵乘法为例,三层循环的最内层每走一步都要从全局内存取数据,同一份数据被反复读取,带宽消耗巨大。切分(split)把一个大张量按维度拆成多个小分块(tile),让每个分块既能装进共享内存或寄存器,又足以支撑成片的计算,从而把“一件一件搬”变成“一车一车搬”。
切分可以沿任意维度进行。设一个三维张量在 M、N、K 方向的总长度分别为 M、N、K,每个分块的大小为 (Tm, Tn, Tk),则分块总数为三个方向分块数的乘积。
调度员在订单图上画出格子,每个格子就是一车货的装载量:格子太小,车次太多,运输成本降不下来;格子太大,货架又吞不下。找到“刚刚好”的格子大小,正是整套 Tiling 策略的起点,也是后面所有调度变换的前提。
- 三个维度$M$、$N$、$K$ 是待切分张量三个方向的总长度,分别对应订单在不同方向上的规模。
- 块尺寸$T_m$、$T_n$、$T_k$ 是每个小分块在各方向上的长度,即每车装载量,受容量约束。
- 向上取整$\lceil \cdot \rceil$ 处理除不尽的情况,避免最后一块越界或漏算。
- 总块数三个方向的分块数相乘,就是需要调度搬运的批次总数。
循环切分变换
分块的尺寸定下来之后,还要给每个格子编上清晰的门牌号。调度员在排班表上写“第 3 车、第 17 件”,把一个大编号拆成了两层:车次号加上车内序号。循环切分(loop split)干的正是这件事。
TVM 中的 T.split 接受一个循环和拆分因子 f,把原来单层循环 i 改写成嵌套的两层:外层 i0 负责遍历分块(第几车),内层 i1 负责遍历块内元素(车内第几个)。两层循环组合起来遍历的范围与原来完全一致,只是索引从一维坐标变成了二维坐标。
设原始循环范围是 $0 \le i < N$,拆分因子为 f,则索引满足下面的关系。
排班表写好后,每个包裹都有了唯一的“车次 + 序号”地址。但有了地址还不够——接下来要决定哪辆车归哪组搬运工,车厢里的包裹按什么顺序卸,这就引出了线程绑定与循环重排。
- 原始索引$i$ 是拆分前循环变量的取值,范围 $0$ 到 $N-1$,相当于包裹的全局编号。
- 外层循环$i_0$ 是拆分后的外层循环变量,表示“第几车”,取值范围由总件数与每车件数决定。
- 内层循环$i_1$ 是拆分后的内层循环变量,表示“车内第几个”,范围严格小于因子 $f$。
- 组合还原$i = i_0 \times f + i_1$ 把二维坐标还原成一维大索引,保证每个元素恰好被访问一次。
线程绑定映射
编号之后是分工。调度员宣布:第 0 组搬运工负责第 0 车,第 1 组负责第 1 车;组内第 k 个搬运工负责车厢里第 k 个格子。于是“车次号”对应上了小组编号,“车内序号”对应上了组员编号——这就是线程绑定(thread binding)。
GPU 的线程是两层结构:线程块(block)负责“组”,线程(thread)负责“组内成员”。TVM 的 T.bind 把切分后的外层循环绑定到 blockIdx(线程块编号),内层循环绑定到 threadIdx(线程编号),从而把循环迭代空间映射到硬件线程空间。每个线程由此算出一个全局数据坐标。
全局坐标的计算遵循一条通用公式,它把块编号、块内规模和块内编号组合成一个唯一的大索引。
一张好的排班表,要让每组搬运工工作量相当、互不抢道。如果一种映射方式让某些线程空转、某些线程排长队,调度员就会用 reorder 调整内外层循环的先后顺序,或者把不同轴换一换绑定方向——这些在 TVM 的调度原语里都是一行代码的事。
- 块编号$\text{blockIdx}$ 表示当前线程块在所有块中的序号,对应第几组搬运工。
- 块内规模$\text{blockDim}$ 表示一个线程块内有多少个线程,即每组搬运工的人数。
- 线程编号$\text{threadIdx}$ 表示线程在块内的序号,对应组内第几个位置。
- 全局编号两者拼接后得到全局坐标,正好对应切分后的某个数据格。
共享内存缓存
在分拣中心,中转货架是吞吐量的胜负手。如果每个打包员都直接跑仓库取件,仓库门口会被挤得水泄不通;统一由小车把一批货卸到货架上,大家再从货架取——这就是共享内存缓存(shared memory)的核心思想。
矩阵乘法里,一个输出分块需要反复使用输入分块的数据。若不经过共享内存,每次累加都穿透到全局内存;若先把输入分块搬进共享内存,块内所有线程都能反复复用这些数据,全局访存次数成倍下降。TVM 中用 storage_scope 与拷贝调度实现这一搬运过程。
共享内存收益的大小可以用“计算访存比”来度量。设分块尺寸为 $T \times T$,为了算出一个输出分块,需要把 A、B 各一块 $T \times T$ 的数据搬进共享内存,再完成 $2T^3$ 次浮点运算。
调度员还要留意货架的“隔板”:共享内存被分成若干 bank,若多个线程同时伸手去拿同一排的包裹,速度会骤降——这就是 bank conflict,常见的对策是在数据间插入 padding。此外,每车货卸完要吹哨统一开工(barrier 同步),防止有人取到半车货就算数。
- 计算量$2T^3$ 是两个 $T \times T$ 输入块相乘得到 $T \times T$ 输出块所需的浮点运算次数。
- 访存量$2T^2$ 是把 A、B 两个分块搬到共享内存的元素总数,每块各 $T^2$ 个。
- 复用比值两者相除得到 $T$,即每搬一个元素可支撑 $T$ 次计算,分块越大越划算。
- 容量上限共享内存容量 $S_{\text{shm}}$ 约束了分块上限,这是公式之外必须遵守的物理边界。
向量化访存
仓库出货时,最有经验的调度员会让叉车把 4 个箱子摞成一个托盘一次运走,而不是来回跑四趟。向量化(vectorization)就是 GPU 世界里的这台叉车。
向量化要求最内层循环的访问地址连续且首地址对齐。TVM 的 T.vectorize 把最内层循环的多次标量访存合并为一条向量访存指令:如果硬件一次能从缓存行搬 16 字节,而每个元素是 4 字节,这一条指令就完成了原本 4 次加载的工作。
向量宽度的计算很简单:用硬件一次传输的字节数除以每个元素的字节数,就得到每条指令能装载的元素个数。
叉车的载货量由硬件决定,但“哪些箱子能拼成一个托盘”却由调度员控制:把最内层循环留给连续地址,让分块的边长取向量宽度的整数倍,向量化就水到渠成,访存带宽也能被真正“喂饱”。
- 缓存行宽度$B_{\text{cache line}}$ 是硬件一次突发传输的字节数,相当于叉车托盘的载货能力。
- 元素位宽$B_{\text{elem}}$ 是每个数据元素占用的字节数,float 为 4 字节,double 为 8 字节。
- 向量宽度$W_{\text{vec}}$ 是每条向量指令能装载的元素个数,即托盘一次能摞多少箱子。
- 指令数$N_{\text{elem}}$ 个元素需要的访存指令数,向量化后指令数降为原来的 $1 / W_{\text{vec}}$。
性能评估调优
方案好不好,不能靠嘴说,要看指标。调度员每天盯着三个数字:打包线空转率、货架翻车次数、通道利用率。GPU 优化里同样有一个关键指标——占用率(occupancy)。
占用率指同时驻留在 GPU 上的活跃线程数与硬件容量的比值。占用率太低,访存延迟无法被隐藏,计算单元只能干等;但占用率也不是越高越好——抬高线程数往往会挤占共享内存和寄存器,导致每个线程“装不下货”。
占用率的计算公式如下,它把每块线程数、驻留块数和设备上限串在一起。
手动调这些参数就像排一个五百人的班表,费时费力。TVM 的 AutoTVM 把分块大小、循环顺序、绑定方式、向量宽度全部视为旋钮,通过反复编译实测,自动搜索最优组合。调度员终于可以不再拍脑袋,而是让数据说话。
- 每块线程数$T_b$ 是每个线程块划分的线程数量,对应每组搬运工的人数。
- 驻留块数$N_b$ 是能同时驻留在 GPU 上的线程块数,受共享内存与寄存器总量约束。
- 硬件上限$N_{\text{max}}$ 是设备支持的最大线程驻留数,是调度方案的天花板。
- 占用率比值两者相除得到占用率,越高越能隐藏访存延迟,但过高会挤占每线程资源。
综合实战案例
把前面所有招式串起来,就是一场完整的搬运演练。以 1024×1024 的矩阵乘法为例:先用 split 把大矩阵切成小块,再 bind 到线程块与线程,reorder 调整遍历顺序,vectorize 收尾,最后在 GPU 上实测对比。
朴素版本每计算一个输出元素,都要从全局内存重新读取 A 和 B,全局访存次数高达 $O(N^3)$ 量级;分块版本把输入块先搬进共享内存,按块反复复用,全局访存次数降为 $O(N^3 / T)$。这就是加速的根本来源。
理论加速比近似等于分块大小 T,因为全局访存次数恰好减少到原来的约 $1/T$。
调度员把整套方案跑完,发现高峰期吞吐量翻了几倍——这就是从理论到工程的分块落地。实际部署中,张量尺寸未必是块的整数倍,最后一个分块会“不满载”,需要用边界判断或掩码修正坐标,保证任何形状的算子都不出错。
- 朴素访存每个输出元素都要在内层循环读 $N$ 次输入,全局访存约 $O(N^3)$ 量级。
- 分块访存每个 $T \times T$ 输入块只从全局读一次进共享内存,全局访存约 $O(N^3 / T)$。
- 理论加速比两者相除得到约 $T$ 倍的访存下降,这是分块收益的上界。
- 实测修正实际加速比用朴素耗时除以分块耗时得到,通常会小于理论值。
AI 应用场景
推荐系统 Embedding 聚合加速
阶段:算子编译与推理优化
解决的问题:推荐模型的 embedding 查询是典型的“大表 + 小访存”:每个用户 ID 都要去全局大表查一行向量。传统实现中每次查询都穿透全局内存,访存时延成为瓶颈。用 Tiling 思路把高频 ID 对应的分块表搬进共享内存,配合向量化合并批量查询——相当于调度员把热门包裹预先挪到中转货架。
具体案例:某广告 CTR 模型在 A100 上,embedding 聚合核函数经 T.split + T.bind + T.vectorize 重构后,端到端推理时延下降 35%。
大模型 Attention 分块计算
阶段:算子开发与编译
解决的问题:Attention 的 QKᵀ 中间矩阵巨大,若完整写回全局内存再读回来,等于让打包员在仓库和货架之间空跑。分块 + 在线 Softmax 让每个分数块驻留共享内存,算完即丢弃,全局访存大幅减少。
具体案例:在 TVM 中实现 FlashAttention 风格的分块注意力,序列长度 4096 时相比标准 Attention 算子吞吐提升 2 倍以上。
边缘端卷积算子部署
阶段:模型部署与端侧适配
解决的问题:Jetson 等边缘设备的共享内存和寄存器更小,卷积必须按通道与空间维度精细切分;忽略 Tiling 的卷积在边缘设备上访存带宽利用率极低,推理耗电高、延迟大。
具体案例:YOLOv5s 的 3×3 卷积经 AutoTVM 搜索到 16×16 输出块后,在 Jetson Xavier NX 上提速 1.8 倍。
Python 演示
import numpy as npimport time def naive_gemm(A, B, C, N): for i in range(N): for j in range(N): for k in range(N): C[i, j] += A[i, k] * B[k, j] def tiled_gemm(A, B, C, N, T): for i0 in range(0, N, T): for j0 in range(0, N, T): b_col = B[:, j0:j0 + T].copy() # 模拟 B 分块搬入共享内存 for k0 in range(0, N, T): a_row = A[i0:i0 + T, k0:k0 + T].copy() # 模拟 A 分块搬入共享内存 b_blk = b_col[k0:k0 + T, :] for ii in range(T): for jj in range(T): acc = 0.0 for kk in range(T): acc += a_row[ii, kk] * b_blk[kk, jj] C[i0 + ii, j0 + jj] += acc def main(): N, T = 64, 8 rng = np.random.default_rng(0) A = rng.random((N, N)) B = rng.random((N, N)) C1 = np.zeros((N, N)) C2 = np.zeros((N, N)) t0 = time.time() naive_gemm(A, B, C1, N) t1 = time.time() tiled_gemm(A, B, C2, N, T) t2 = time.time() diff = np.max(np.abs(C1 - C2)) speed = (t1 - t0) / max(t2 - t1, 1e-9) print("naive=%.4fs tiled=%.4fs diff=%.2e speedup=%.1fx" % (t1 - t0, t2 - t1, diff, speed)) if __name__ == "__main__": main()
for i0 in range(0, N, T):外层循环按分块大小 T 切分行方向,相当于决定第一批运输车次。b_col = B[:, j0:j0 + T].copy()先把 B 的第 j0 列分块“搬运”出来,模拟写入共享内存,供后续多轮复用。for k0 in range(0, N, T):在第三个维度继续切分,形成三维 Tiling 的调度表。a_row = A[i0:i0 + T, k0:k0 + T].copy()A 的分块同样搬入“共享内存副本”,避免每层循环都读原数组。acc = 0.0累加器留在“寄存器”中,不频繁落内存,这是分块能减少访存的关键。diff = np.max(np.abs(C1 - C2))校验两种写法结果一致,确保切分不改变计算语义。
常见误区
误区分块肯定是越大越好呀,公式里都写了计算访存比等于 T,T 越大访存越少,那我直接把 T 设成 512,性能一定起飞。
正解公式只描述了复用收益,却忽略了约束:共享内存和寄存器都有限,T 太大会导致编译失败、线程块无法驻留,甚至寄存器溢出(spill)把数据倒回本地内存,反而更慢。正确的做法是在容量约束下,寻找占用率与复用率平衡的 T,再用 AutoTVM 实测确认。
为何常见:初学者看到“计算访存比 = T”就以为单调递增,忽视了硬件容量和驻留约束这两个前提。
误区我用 T.split 把循环写成两层了,为什么性能一点变化都没有?切分不是能提升性能吗?
正解split 只是把数据重新编号,并未改变访问顺序,也没让数据靠近计算。真正起作用的是配对的 bind(把编号映射到线程)、reorder(调整搬运顺序)、vectorize(合并访存)以及共享内存搬运指令。只 split 不配套,编译器甚至可能把两层循环优化回原样,等于白切。
为何常见:把“循环变换”误当成“数据调度”,以为改写索引本身就有收益,忽略了调度变换必须组合使用。
误区只要我把数据 copy 到共享内存,就一定比直接读全局内存快。
正解共享内存只是把“远路运输”变成“近路取货”,前提是这份数据要被多次复用。如果一份数据只用一次,搬运本身就是多余的开销;如果多个线程同时访问同一 bank 的不同地址触发 bank conflict,近路也会堵车;同步屏障还会引入等待。共享内存适合“一搬多用”的场景,而不是无脑搬运。
为何常见:只看到共享内存时延低,忽略了搬运成本、冲突开销和同步等待这三笔隐形账。
误区反正我们模型里矩阵尺寸都是 512 的倍数,边界那块不用管。
正解工程中算子输入形状千变万化,当 N 不能被 T 整除时,最后一个分块必然越界或不满载。忽略边界会直接算错甚至访问非法内存;常见做法是用条件判断、掩码或让尾部线程提前退出,保证正确性后再谈性能。benchmark 里的规整形状只是为了方便,不代表真实世界。
为何常见:教程和基准测试多用整齐的幂次尺寸,掩盖了余数处理这个必修课。
面试题
Q1. TVM 中 Tiling 分块的核心目标是什么?
- A.让数据尽量贴近高速缓存复用← 正确答案
- B.把浮点数换成整数来加速
- C.减少代码行数以降低编译时间
- D.让张量形状变成正方形
答案:A。分块的核心是让大张量中相互复用的数据尽量落在共享内存或寄存器里,减少对全局内存的远距离访存——就像把热门包裹提前挪到中转货架。B 是数值精度问题,C 是编译期问题,D 只是形状,均与访存优化无关。
Q2. GPU 存储层次中延迟最低的是哪个?
- A.全局内存 Global Memory
- B.共享内存 Shared Memory
- C.寄存器 Registers← 正确答案
- D.L2 缓存 L2 Cache
答案:C。寄存器是每个线程私有的“贴身托盘”,访问几乎零延迟;共享内存同块线程共用但需要指令搬运;全局内存容量最大、延迟数百周期。记住这条“远近顺序”是理解 Tiling 的前提。
Q3. 在 TVM 中,T.split 的作用是什么?
- A.将一个循环拆成内外两层嵌套循环← 正确答案
- B.把一个算子拆成两个等价算子
- C.将两个循环合并成一个循环
- D.将数据从全局搬到共享内存
答案:A。T.split 把单层循环变成“第几车 + 车内第几个”两层嵌套循环,等价于 i = i0 × f + i1。B 是算子拆分,C 是 fuse 的功能,D 是存储调度,均不是 split 的职责。
Q4. 线程绑定中,全局索引的正确计算公式是?
- A.blockIdx × blockDim + threadIdx← 正确答案
- B.blockIdx + blockDim ÷ threadIdx
- C.blockIdx − blockDim + threadIdx
- D.blockIdx ÷ blockDim + threadIdx
答案:A。全局索引 = blockIdx × blockDim + threadIdx,与“车次 × 每车人数 + 车内序号”完全同构,保证每个数据格恰好由一个线程负责。其余选项的运算符都会导致数据重叠或遗漏。
Q5. 共享内存发生 bank conflict 时会发生什么?
- A.同 Bank 请求被硬件串行化执行← 正确答案
- B.数据自动复制到全局内存
- C.线程自动重启重新计算
- D.编译器跳过共享内存读写
答案:A。共享内存由多个 bank 组成,多个线程同时访问同一 bank 的不同地址时,硬件只能把这些访问串行化执行,性能骤降。数据不会转存全局,线程也不会重启,编译器更不会跳过读写。
Q6. 向量化加载要求数据满足什么条件?
- A.地址连续且首地址按对齐← 正确答案
- B.数据必须是无符号整数
- C.数组必须保证一维连续存储
- D.线程总数必须是奇数个
答案:A。向量加载依赖“叉车一次搬一摞箱子”:箱子必须连续摆放(地址连续)且方向对齐,才能合并成一条指令。元素类型、存储维度与线程数的奇偶都与合并条件无关。
Q7. 分块过大会导致什么后果?
- A.共享内存超限或寄存器溢出← 正确答案
- B.数组越界但仍能算对
- C.浮点精度自动提升
- D.算子自动变成卷积实现
答案:A。分块越大,需要的共享内存和寄存器越多,超出后线程块无法驻留或发生寄存器溢出,性能不升反降。越界问题与“是否能算对”无关,浮点精度不会自动提升,算子类型更不会改变。
Q8. 下列哪些属于 AutoTVM 可搜索的调度参数?
- A.循环切分的分块尺寸大小← 正确答案
- B.内外层循环的执行顺序← 正确答案
- C.是否绑定到线程块维度← 正确答案
- D.模型训练的迭代轮数
答案:ABC。AutoTVM 搜索的正是调度空间:分块大小、循环顺序、是否绑定、是否向量化都属于可配置项。而训练轮数属于模型训练流程,与算子编译调度无关,不在搜索范围之内。
Q9. 使用共享内存加速需要付出哪些额外开销?
- A.从全局搬运到共享的额外成本← 正确答案
- B.线程间同步屏障的等待开销← 正确答案
- C.Bank 冲突导致的串行访问← 正确答案
- D.全局内存容量自动减半
答案:ABC。共享内存不是免费的:搬运有成本、同步有等待、bank 冲突会串行化。只有复用次数足够多时收益才大于开销。D 是硬件属性,不会随使用共享内存而变化。
Q10. 关于边界分块的处理方式,下列说法正确的有?
- A.需要判断余数并修正坐标← 正确答案
- B.使用边界掩码避免越界访问← 正确答案
- C.让尾部线程提前退出循环← 正确答案
- D.任何情况下都能忽略余数
答案:ABC。真实张量不能保证被整除,必须处理余数:坐标判断、掩码、提前退出都是合法手段,直接忽略会越界。D 的说法在任何情况下都不成立。