共享内存体位与冲突

优化内存吞吐量 (throughput)不仅仅是把数据移到最快的可用存储中。一旦数据存入共享内存（SMem），线程访问这些数据的方式就决定了内核是能达到峰值带宽，还是因序列化而停滞。NVIDIA GPU架构中的共享内存被分成大小相等的内存模块，称为体位（banks），它们可以同时被访问。当单个warp中的多个线程尝试读写同一体位内不同地址时，就会发生体位冲突，迫使硬件按顺序重放内存请求。

共享内存体位的架构

共享内存组织成32个连续的体位，对应一个warp中的32个线程。每个体位在每个时钟周期具有32比特（4字节）的带宽。内存地址到体位的映射由地址模32决定。具体来说，对于标准的32比特字长，连续的字被分配到连续的体位。

如果一个warp执行加载指令，硬件会检查活动线程请求的地址。理想情况是无冲突访问，即每个线程引用一个独特的体位。这使得内存控制器能够在一个事务中处理所有32个请求。

然而，深度学习 (deep learning)工作负载通常涉及矩阵乘法和卷积，这些操作需要访问多维数组。根据步长和张量形状，这些访问模式可能无意中对齐 (alignment)到同一个体位。

给定32比特地址的体位索引计算如下：

$\text{体位索引} = \left( \frac{\text{字节地址}}{4} \right) \pmod{32}$

当多个线程映射到同一个体位索引时，硬件将请求拆分为$n$个单独的事务，其中$n$是任何单个体位中冲突的最大数量。这意味着一个$n$向体位冲突会将有效共享内存带宽降低$n$倍。

体位访问模式与冲突可视化。左侧显示线性访问模式，其中每个线程映射到唯一的体位（无冲突）。右侧显示步进访问模式（步长为2），导致2向体位冲突，因为两个线程同时映射到同一个体位。

识别张量布局中的冲突

在张量处理中，数据通常以行主序格式存储。按行访问数据块通常会产生单位步长访问，即无冲突。然而，访问数据块内的列常常会导致步进访问。

考虑一个大小为$32 \times 32$的块，其中存储float元素。两个垂直元素之间的步长是32个字。如果一个warp尝试加载一列（例如，线程$i$加载$A[i][0]$），每个线程访问的地址相隔正好32个字。由于$32 \pmod{32} = 0$，warp中的所有32个线程都映射到体位0。这会导致32向体位冲突，使执行序列化，并将吞吐量 (throughput)降低到峰值能力的1/32。

这种情况在GEMM（通用矩阵乘法）操作中很常见，比如当一个矩阵被转置时，或者在为Tensor Core加载数据块时。

优化策略：填充

解决体位冲突的最简单方法是填充。通过在分配中添加一个“虚拟”列，物理步长会改变，而逻辑访问模式保持不变。对于一个$32 \times 32$的块，将其分配为$32 \times 33$会将步长从32变为33。

当步长为33时：

1. 线程0访问索引0 $\to$ 体位0。
2. 线程1访问索引33 $\to$ 体位1 ($33 \pmod{32} = 1$)。
3. 线程$k$访问索引$33k \to$ 体位 $k \pmod{32}$。

由于$\text{gcd}(33, 32) = 1$，访问会循环遍历所有体位，完全消除冲突。尽管有效，但填充会增加共享内存的使用量，而共享内存是一种稀缺资源。在由共享内存容量限制占用率的高度优化的内核中，浪费3%的空间（32列中浪费1列）可能很可观。

优化策略：交织

为避免填充带来的内存开销，现代编译器（如TVM和Triton内部的逻辑）采用地址交织（通常称为XOR交织）。交织在不改变分配大小的情况下置换元素的存储位置。

这个思路是使用按位XOR操作根据行索引修改列索引。一个常见的2D块交织模式是：

$\text{列}_{\text{物理}} = \text{列}_{\text{逻辑}} \oplus \text{行}_{\text{逻辑}}$

或者，对于更宽的块：

$\text{列}_{\text{物理}} = \text{列}_{\text{逻辑}} \oplus \left( \frac{\text{行}_{\text{逻辑}}}{\text{块宽度}} \right)$

通过将行位异或到列地址中，每行的数据在体位之间循环移位。当线程沿列读取（增加行索引）时，物理列索引会移位，从而将访问分散到不同的体位。

这项技术使得编译器能够保持紧凑的存储，同时确保行主序和列主序访问模式都具有减少或消除的体位冲突。

不同内存访问策略的吞吐量 (throughput)比较。该可视化突出显示了N向冲突的严重性能损失，以及使用填充或交织技术恢复性能的情况。

向量 (vector)化加载和体位宽度

高级编译策略还必须考虑向量化 (quantization)指令，例如LDS.128（一次加载128比特或4个浮点数）。虽然基本体位宽度是32比特，但硬件可以将体位组合起来以处理更宽的加载。

然而，向量化会增加步长要求。如果每个线程加载128比特，它会消耗来自4个连续体位的带宽。为了使一个warp无冲突地执行此操作，线程访问之间的步长必须有效跳过128比特（或与超体位结构完美对齐 (alignment)）。在成本模型中模拟此行为的编译器必须确保向量化不会重新引入冲突，这些冲突的弊大于更少指令发布的好处。

例如，MLIR的向量分发pass通常包含特定的重写规则，以检查共享内存中张量的最内层维度是否连续，以及维度大小是否是向量宽度的倍数。如果满足这些条件，编译器会发出向量加载；否则，它会回退到标量加载，以避免未对齐体位访问的开销。