优化 global memory 访问¶
一个线程读取一行中的多个元素时，有四种写法。本页给出它们在两个 workload 上的实测对比，以及选择的依据。
确认 DRAM 带宽是否为当前的瓶颈¶
Elementwise 与 Reduction 是典型的访存受限 kernel。本页每条建议都写明触发条件、成因，以及反例与正例代码。
本页的实测都在同一组条件下取得：输入大于 L2 的 60 MiB，且 block 数足以填满整卡（H200 有 132 个 SM）。此时 DRAM 带宽是主要瓶颈，访存模式的差别直接反映在性能上。
适用范围
在这组条件之外，主要瓶颈可能是其他因素，本页的部分结论会反转。
下表按这两个条件划分三个区间，逐行给出判据，以及 global memory 与 shared memory 两页的结论在各区间的用法。表中列出的是主导瓶颈；kernel 越复杂，同时起作用的因素越多。
| 区间 | 判据 | 主要瓶颈 | 结论的用法 |
|---|---|---|---|
| 带宽饱和 | 输入 > 60 MiB，block 数在 SM 数的两倍以上 | DRAM 带宽，即 sector 利用率 | 直接适用 |
| 数据小 | 输入装得进 L2，单次耗时在几十微秒以内 | kernel 发射的固定开销、缓存状态 | 避开反例即可，换 access pattern 没有收益 |
| block 少 | block 数不到 SM 数的两倍 | 每条载入指令的宽度、在飞的字节数 | 优先保证载入宽度，每项改动单独实测 |
- 数据小时，发射开销与缓存状态占主导。 同一个行求和 kernel 的四种 access pattern（fp16，256 线程，时钟未锁）在 65536 × 4096（512 MB）上测得的访存带宽是 4.20 到 4.43 TB/s，彼此相差不超过 6%；换成 2048 × 4096（16 MB，装得进 L2）后单次耗时十几微秒，同一个 access pattern 两次测量之间可差三倍。在这个区间里，制约性能的是其他因素，更换 access pattern 没有收益。
- block 少时，载入宽度带来的收益大于合并规则算出的差别。 warp 数量不足时，无法靠并发的请求数掩盖访存延迟，只能让每个请求更宽、每个线程持有更多在飞的字节。牺牲载入宽度来换取其他收益的改法，在这个区间都可能得到相反的结果。
合并 global memory 的访存¶
一次 global memory 读取的数据访问单位是 32 字节的 sector。
| 名称 | 大小 | 是什么 |
|---|---|---|
| cache line | 128 字节 | L1 与 L2 的缓存行，也是缓存查找的单位 |
| sector | 32 字节 | 一条 cache line 由 4 个 sector 组成，L1 与 L2 之间按 sector 传输 |
缓存查找时以 cache line 为单位，搬运数据时以 sector 为单位：某个 sector 未命中，L1 就只向 L2 请求这一个 sector，不必把整条 cache line 都拉过来。因此取 1 个字节和取满 32 个字节的代价相同，一条访存指令的效率由 sector 利用率衡量：真正用到的字节 / (覆盖的 sector 数 × 32)。
硬件把一个 warp 的 32 个访问合并成尽可能少的 32 字节事务。事务数降到最少需要同时满足三个条件：
- 地址连续 —— 同一条指令里 32 个线程的地址首尾相接，不留空洞；
- 按 32 字节对齐 —— 起始地址是 32 的倍数，一段数据不会多占一个 sector；
- 每个线程一次取满 16 字节 —— 一条指令覆盖 \(32 \times 16 = 512\) 个连续字节，即 16 个满载的 sector。
三个条件同时成立时，这条访存指令对硬件最友好。
一个线程要读 \(V\) 个元素时（\(V\) = 一行的元素数 / 线程数），有四种 access pattern。
blocked —— 每个线程负责一段连续的元素。固定 c 时相邻线程的地址相隔 \(V\) 个元素，违反第一个条件，sector 利用率是 \(1/V\)：
striped —— 相邻线程取相邻元素。地址连续，但每个线程一次只取一个元素，违反第三个条件，\(V\) 个元素需要 \(V\) 条指令：
blocked + 向量化 —— 每个线程仍负责一段连续的元素，但改用 T.vectorized 一次读满 16 字节，三个条件全部满足：
buf = T.alloc_local((V,), dtype)
for c in T.vectorized(V):
buf[c] = X[row, tx * V + c]
for c in T.serial(V):
acc[0] = acc[0] * buf[c]
staged —— 由 T.Parallel 完成搬运，消费阶段改从 shared memory 读取，同样满足三个条件：
sh = T.alloc_shared((threads, V + pad), dtype)
for t, c in T.Parallel(threads, V):
sh[t, c] = X[row, t * V + c]
T.sync_threads()
for c in T.serial(V):
acc[0] = acc[0] * sh[tx, c]
四种 access pattern 的差别在于：「哪个线程读哪些元素、一次读多宽」这个映射由谁决定。
T.serial 的语义是循环体由单个线程顺序执行，索引表达式被逐字翻译成访存指令，不做合并，也不做向量化，因此编程者写出的模式就是硬件看到的模式。
T.vectorized、T.Parallel、T.copy 的映射则由 TileLang 的 layout inference 决定。三者的差别在于编程者需要写明的内容：
T.vectorized写明每个线程一次访问的宽度，线程映射由 layout inference 推导；T.Parallel不必写宽度，循环维度如何分给线程、一次读多宽都由 layout inference 决定；T.copy只写源区域与目标区域，整段搬运由它生成；需要自行指定推导结果时，另有coalesced_width与loop_layout两个参数。
其余的向量化、地址对齐，以及在 shared memory 一侧避开 bank 冲突，都由 layout inference 负责。这些正是对硬件友好、但手写容易出错的部分。
用 T.serial 手写下标时，三个条件由编程者逐一保证；交给 layout inference 时，只需写明搬运的范围。 各种写法的选择与各自达到的带宽见下文的实测。
实测对比¶
本节在 H200 上实测两个 workload，比较四种 access pattern 各自达到的访存带宽（搬运的字节数除以 kernel 耗时，单位 TB/s）。两个 workload 的计算对元素处理顺序的要求不同，这一要求决定了哪几种 access pattern 可用。
测试中 SM 时钟锁在 1830 MHz；输入 bf16 的 \(65536 \times 4096\)（512 MB，必须大于 L2 的 60 MiB，否则测到的是 L2 带宽）；每个配置跑三次，三次的结果一致到 ±0.5%。staged 在表里占两列：一列不加 pad（此时 stride 恰好是 \(V\) 个 word，产生 bank 冲突，见优化 shared memory 访问），一列是在若干个 pad 取值中测到的最优值。
workload 1，按行求乘积：一行的元素相乘，先后顺序不影响结果，每行只输出一个值，四种 access pattern 都可用。
| 线程数 | \(V\) | blocked TB/s |
striped TB/s |
blocked + 向量化 TB/s |
staged 不加 pad TB/s |
staged 加 pad TB/s |
|---|---|---|---|---|---|---|
| 512 | 8 | 3.02 | 3.04 | 3.22 | 3.05 | 3.02 |
| 256 | 16 | 1.83 | 3.31 | 3.81 | 3.79 | 3.79 |
| 128 | 32 | 0.95 | 3.36 | 3.99 | 3.81 | 3.98 |
| 64 | 64 | 0.48 | 3.49 | 3.32 | 2.80 | 4.02 |
| 32 | 128 | 0.47 | 3.43 | 3.11 | 2.83 | 3.88 |
workload 2，按行求串行前缀积：每个位置的结果依赖它左边的全部元素，顺序不能变，整行都要写回。striped 在这里不可用（线程无法持有连续的一段）。
| 线程数 | \(V\) | blocked TB/s |
blocked + 向量化 TB/s |
staged 不加 pad TB/s |
staged 加 pad TB/s |
|---|---|---|---|---|---|
| 512 | 8 | 0.92 | 4.20 | 2.47 | 3.15 |
| 256 | 16 | 0.42 | 3.85 | 1.63 | 3.29 |
| 128 | 32 | 0.34 | 2.76 | 0.89 | 3.35 |
| 64 | 64 | 0.26 | 1.80 | 0.46 | 3.69 |
| 32 | 128 | 0.27 | 1.60 | 0.46 | 3.24 |
access pattern 的取舍¶
-
逐元素的 blocked 在 \(V > 1\) 时总是最差的 access pattern。 固定
c时相邻线程的地址相隔 \(V\) 个元素，sector 利用率是 \(1/V\)，所以 \(V\) 越大越差：workload 1 的表中从 \(V = 8\) 的 3.02 掉到 \(V = 64\) 的 0.48。这个关系由访存合并的规则决定，不随形状改变。 -
\(V\) 小时用向量化的 blocked；\(V\) 大到寄存器压力压低占用率时，改用加了 pad 的 staged。 向量化把整段留在寄存器里（bf16 是每线程 \(V/2\) 个），staged 把它放进 shared memory，以一次同步为代价节省寄存器。翻转点取决于 kernel 其余部分剩下的寄存器预算，没有固定的 \(V\)：上面两个 workload 在同一个行宽下分别落在 \(V = 64\) 与 \(V = 32\)。翻转点需要在具体的 kernel 上实测。
-
staged 的 shared 缓冲要避开 bank 冲突。 声明成
(threads, V)时 stride 恰好是 \(V\) 个 word，\(V\) 为 2 的幂就一定产生冲突；pad 的算法与候选见优化 shared memory 访问。workload 2 的表里，同一个配置不加 pad 是 0.46，加 pad 是 3.69。 -
striped 完全合并，但每个元素需要一条指令。 因此它好于逐元素的 blocked、差于向量化的 blocked（\(V = 16\) 上 3.31 对 1.83 与 3.81），适合改动量比最后一点带宽更重要的场合。它使每个线程持有的元素不连续，因此不能用于要求线程持有连续一段元素的计算，例如串行前缀积。
下面两段是推荐 access pattern 的完整模板，M、N、V、threads、pad、dtype 都是编译期常量。「合并 global memory 的访存」一节的四段代码中，逐元素的 blocked 是反例；striped 可用，但不是最快的一种（见取舍的第 4 条）。
推荐的 access pattern，小 \(V\) —— 向量化的 blocked：
@T.prim_func
def main(X: T.Tensor((M, N), dtype), Out: T.Tensor((M, threads), "float32")):
with T.Kernel(M, threads=threads) as row:
tx = T.get_thread_binding()
buf = T.alloc_local((V,), dtype)
acc = T.alloc_local((1,), "float32")
acc[0] = T.cast(1.0, "float32")
for c in T.vectorized(V): # 一条 16 字节向量读
buf[c] = X[row, tx * V + c]
for c in T.serial(V): # 消费，顺序随计算需要
acc[0] = acc[0] * T.cast(buf[c], "float32")
Out[row, tx] = acc[0]
推荐的 access pattern，大 \(V\) —— 加了 pad 的 staged（翻转点见取舍的第 2 条）：
@T.prim_func
def main(X: T.Tensor((M, N), dtype), Out: T.Tensor((M, threads), "float32")):
with T.Kernel(M, threads=threads) as row:
tx = T.get_thread_binding()
sh = T.alloc_shared((threads, V + pad), dtype) # pad 的选法见 shared memory 那一篇
acc = T.alloc_local((1,), "float32")
acc[0] = T.cast(1.0, "float32")
for t, c in T.Parallel(threads, V): # 搬运，映射由 layout inference 给出
sh[t, c] = X[row, t * V + c]
T.sync_threads()
for c in T.serial(V): # 消费
acc[0] = acc[0] * T.cast(sh[tx, c], "float32")
Out[row, tx] = acc[0]