CUDA Shared Memory 在向量化指令下的访存机制

更新日志

🔥2024.03.24 添加了关于对 128 Bit 下的 Bank Conflict 的讨论
🔥2024.04.01 修正了 128 Bit 下第五个例子的错误代码，感谢知乎用户 @Alan小分享指证～

NVIDIA GPU 上的内存结构从慢到快分为 Global Memory、L2 缓存、L1TEX 缓存/Shared Memory 和寄存器，从 Volta/Turning 开始其中 L1TEX 缓存和 Shared Memory 在物理上放在了同一块芯片上，拥有相似的延时和带宽¹，因此如何掌握 Shared Memory 对性能而言就变得尤为重要。可惜的是，NVIDIA 官方的 CUDA 编程手册中介绍 Shared Memory 的教程其实只介绍了在每个线程访问一个 4 字节（32 位宽）的元素时 Bank Conflict 和广播机制，但对使用常用的向量化访存指令如 LDS.64 或 LDS.128 这种一次能访问 8 个字节（64 位宽）或 16 个字节（128 位宽）元素的情况却鲜有资料讨论。这篇文章大概就是想结合网络上的一些讨论以及通过 Microbenchmark 对这些细节来一探究竟。需要注意的是这篇文章的结论仅在 Turing 架构的 GPU 上验证过，其他架构的 GPU 可能会产生变化（欢迎评论区交流🤪）。

Shared Memory 模型

我们可以把 Shared Memory 它看作是一个长度为 $N$ 的数组，每个数组元素的大小是 4 字节（比如可以是一个 int 或 float），这个数组对于同一个 Thread Block 中的所有线程都是可见的，但不同 Thread Block 之间的 Shared Memory 不能互相访问。其中 Shared Memory 最值得注意的点机制是其本身被划分称了 32 个 Bank，其中数组中第 $i$ 个元素存储在了第 $i \bmod 32$ 个 Bank 上。Shared Memory 访存机制可以总结为如下两条：

如果一个 Warp 中的两个不同的线程同时访问了同一个元素，那么就会触发广播机制，即可以合并成对 Shared Memory 的一次访问；
如果一个 Warp 中两个不同的线程同时访问了同一个 Bank 中的不同元素，那么就会产生Bank Conflict，可以理解成一个 Bank 同一时间只能产生吞吐一个元素，因此这两个线程的访存请求会被串行，因而会影响性能；

在 Nsight Compute 上，我们可以通过 Shared Memory 上的 Wavefront 数目来理解 Shared Memory 访存性能，Wavefront 越多说明访存需要的时间越长。

下面几张图片举了几个例子方便理解：

只会触发广播机制，没有 Bank Conflict，需要 1 个 Wavefront

不会触发广播机制，没有 Bank Conflict，需要 1 个 Wavefront

既会触发广播机制，也有 Bank Conflict，需要 4 个 Wavefront（注意第 18 个 Bank）

向量化访存指令

前面在讨论 Shared Memory 上的访存时，我们的 Shared Memory 模型只讨论了一个 Warp 内每个线程所访问的元素。在涉及到向量化访存时这样的模型就不起效果了，因为通过一个 LDS.64 或 LDS.128 指令就可以让一个线程一次性访问 8 个或 16 个字节（相当于 2 个或 4 个元素）。

正确的做法应该是就每个 Wrap 内所产生的每个 Memory Transaction而非每个 Warp 或每条指令来讨论（参考这里）。那么一个在 Shared Memory 上的向量化指令 LDS.64 或 LDS.128 指令到底对应多少个 Memory Transaction？我并没有找到 NVIDIA 给出的官方答案，通过一些网络上的讨论和我自己的 Microbenchmark，我对结合了向量化指令的 Shared Memory 的访存机制的推测如下。

首先一个原则是一个 Warp 中所有线程在同时执行一条 Shared Memory 访存指令时会对应到 1 个或多个 Memory Transaction，一个 Memory Transaction 最长是 128 字节。如果一个 Warp 内在同一时刻所需要的访存超过了 128 字节，那么会则被拆成多个 Transaction 进行。因为一个 Warp 同一时刻执行的访存指令的位宽应该是一样的（即例如不存在线程 0 执行 LDS.32 而线程 1 执行 LDS.128），因此我们只需要对 64 位宽和 128 位宽的访存指令分别讨论即可。

64 位宽的访存指令

对于 64 位宽的访存指令而言，除非触发广播机制，否则一个 Warp 中有多少个活跃的 Half-Warp 就需要多少个 Memory Transaction，一个 Half-Warp 活跃的定义是这个 Half-Warp 内有任意一个线程活跃。触发广播机制只需满足以下条件中的至少一个：

对于 Warp 内所有活跃的第 $i$ 号线程，第 $i~\mathrm{xor}~1$ 号线程不活跃或者访存地址和其一致；
对于 Warp 内所有活跃的第 $i$ 号线程，第 $i~\mathrm{xor}~2$ 号线程不活跃或者访存地址和其一致；

如果触发了广播机制，那么两个 Half-Warp 内的 Memory Transaction 可以合并成一个。

我们看几个例子：

Case 1: 活跃线程全部在第 1 个 Half-Warp 内，需要 1 个 Memory Transaction，没有 Bank Conflict，需要 1 个 Wavefront

Case 2: 活跃线程分散在了 2 个 Half-Warp 内，需要 2 个 Memory Transaction，没有 Bank Conflict，需要 2 个 Wavefront（注意第 15 号和第 16 号线程）

Case 3: 活跃线程分散在了 2 个 Half-Warp 内，但因为触发了广播机制中的第一条，因此仍然只需要 1 个 Memory Transaction，没有 Bank Conflict，需要 1 个 Wavefront

Case 4: 活跃线程分散在了 2 个 Half-Warp 内，看似好像触发了广播机制，但其实并没有，因为第一个 Half-Warp 触发的是第一条，第二个 Half-Warp 触发的是第二条，因此仍然需要 2 个 Memory Transaction，没有 Bank Conflict，需要 2 个 Wavefront

Case 5: 活跃线程分散在了 2 个 Half-Warp 内，没有触发广播机制，需要 2 个 Memory Transaction，没有 Bank Conflict，需要 2 个 Wavefront

可以通过 Nsight Compute 跑一跑下面的代码并观察和 Shared Memory 的相关 Metric 来验证上面这五个例子：

smem_64bit.cu#include <cstdint>

__global__ void smem_1(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  if (tid < 16) {
    reinterpret_cast<uint2 *>(a)[tid] =
        reinterpret_cast<const uint2 *>(smem)[tid];
  }
}

__global__ void smem_2(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  if (tid < 15 || tid == 16) {
    reinterpret_cast<uint2 *>(a)[tid] =
        reinterpret_cast<const uint2 *>(smem)[tid == 16 ? 15 : tid];
  }
}

__global__ void smem_3(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  reinterpret_cast<uint2 *>(a)[tid] =
      reinterpret_cast<const uint2 *>(smem)[tid / 2];
}

__global__ void smem_4(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  uint32_t addr;
  if (tid < 16) {
    addr = tid / 2;
  } else {
    addr = (tid / 4) * 4 + (tid % 4) % 2;
  }
  reinterpret_cast<uint2 *>(a)[tid] =
      reinterpret_cast<const uint2 *>(smem)[addr];
}

__global__ void smem_5(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  reinterpret_cast<uint2 *>(a)[tid] =
      reinterpret_cast<const uint2 *>(smem)[tid % 16];
}

int main() {
  uint32_t *d_a;
  cudaMalloc(&d_a, sizeof(uint32_t) * 128);
  smem_1<<<1, 32>>>(d_a);
  smem_2<<<1, 32>>>(d_a);
  smem_3<<<1, 32>>>(d_a);
  smem_4<<<1, 32>>>(d_a);
  smem_5<<<1, 32>>>(d_a);
  cudaFree(d_a);
  cudaDeviceSynchronize();
  return 0;
}

上面的例子中并没有列举出有 Bank Conflict 的情况，那么 Bank Conflict 在这种情况下应该如何考虑呢？正如前面提到的那样，我们只需要计算每个 Memory Transaction 中的 Bank Conflict 数目然后加起来就好了（因为 Memory Transaction 是串行的）。

128 位宽的访存指令

128 位宽的访存指令和 64 位宽的访存指令是类似的，不同的是需要以 Half-Warp 为单位来计算，对于每个 Half-Warp 而言，除非触发广播机制，这个 Half-Warp 中有多少个活跃的 Quarter-Warp 就需要多少个 Memory Transaction，一个 Quarter-Warp 活跃的定义是这个 Quarter-Warp 内有任意一个线程活跃。类似地，如果触发广播机制那么两个 Quarter-Warp 中的 Transaction 就可以被合并成一个。触发广播机制的条件和 64 位宽的访存指令是一样的（注意广播机制是以整个 Warp 为单位考虑）。这也就意味着假设一个 Warp 中 32 个线程都活跃，即使它们的访存地址都一样，也需要 2 个 Memory Transaction。

同样来看几个例子：

Case 1: 活跃线程分散在了 2 个 Half-Warp 和 2 个 Quarter-Warp 内，每个 Half-Warp 需要 1 个 Memory Transaction，总共需要 2 个 Memory Transaction，没有 Bank Conflict，需要 2 个 Wavefront

Case 2: 活跃线程分散在了 1 个 Half-Warp 和 2 个 Quarter-Warp 内，需要 1 个 Memory Transaction，没有 Bank Conflict，需要 1 个 Wavefront

Case 3: 活跃线程分散在了 2 个 Half-Warp 和 4 个 Quarter-Warp 内，但触发了广播机制（第一条），每个 Half-Warp 需要 1 个 Memory Transaction，总共需要 2 个 Memory Transaction，没有 Bank Conflict，需要 2 个 Wavefront

Case 4: 活跃线程分散在了 2 个 Half-Warp 和 4 个 Quarter-Warp 内，没有触发广播机制，每个 Half-Warp 需要 2 个 Memory Transaction，总共需要 4 个 Memory Transaction，没有 Bank Conflict，需要 4 个 Wavefront

Case 5: 活跃线程分散在了 2 个 Half-Warp 和 4 个 Quarter-Warp 内，但触发了广播机制（第一条和第二条），每个 Half-Warp 需要 1 个 Memory Transaction，总共需要 2 个 Memory Transaction，但因为每个 Memory Transaction 中有 1 个 Bank Conflict，因此会拆分成 4 个 Memory Transaction，对应需要 4 个 Wavefront

Case 6: 活跃线程分散在了 2 个 Half-Warp 和 4 个 Quarter-Warp 内，没有触发广播机制，每个 Half-Warp 需要 2 个 Memory Transaction，总共需要 4 个 Memory Transaction，没有 Bank Conflict，需要 4 个 Wavefront

同样可以通过 Nsight Compute 跑一跑下面的代码并观察和 Shared Memory 的相关 Metric 来验证上面这几个例子：

smem_128bit.cu#include <cstdint>

__global__ void smem_1(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  if (tid == 15 || tid == 16) {
    reinterpret_cast<uint4 *>(a)[tid] =
        reinterpret_cast<const uint4 *>(smem)[4];
  }
}

__global__ void smem_2(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  if (tid == 0 || tid == 15) {
    reinterpret_cast<uint4 *>(a)[tid] =
        reinterpret_cast<const uint4 *>(smem)[4];
  }
}

__global__ void smem_3(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  reinterpret_cast<uint4 *>(a)[tid] = reinterpret_cast<const uint4 *>(
      smem)[(tid / 8) * 2 + ((tid % 8) / 2) % 2];
}

__global__ void smem_4(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  uint32_t addr;
  if (tid < 16) {
    addr = (tid / 8) * 2 + ((tid % 8) / 2) % 2;
  } else {
    addr = (tid / 8) * 2 + ((tid % 8) % 2);
  }
  reinterpret_cast<uint4 *>(a)[tid] =
      reinterpret_cast<const uint4 *>(smem)[addr];
}

__global__ void smem_5(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  reinterpret_cast<uint4 *>(a)[tid] =
      reinterpret_cast<const uint4 *>(smem)[(tid / 16) * 4 + (tid % 16) / 8 + (tid % 8) / 4 * 8];
}

__global__ void smem_6(uint32_t *a) {
  __shared__ uint32_t smem[128];
  uint32_t tid = threadIdx.x;
  for (int i = 0; i < 4; i++) {
    smem[i * 32 + tid] = tid;
  }
  __syncthreads();
  uint32_t addr = (tid / 16) * 4 + (tid % 16 / 8) * 8;
  if (tid < 16) {
    addr += (tid % 4 / 2) * 2;
  } else {
    addr += (tid % 4 % 2) * 2;
  }
  reinterpret_cast<uint4 *>(a)[tid] =
      reinterpret_cast<const uint4 *>(smem)[addr];
}

int main() {
  uint32_t *d_a;
  cudaMalloc(&d_a, sizeof(uint32_t) * 128);
  smem_1<<<1, 32>>>(d_a);
  smem_2<<<1, 32>>>(d_a);
  smem_3<<<1, 32>>>(d_a);
  smem_4<<<1, 32>>>(d_a);
  smem_5<<<1, 32>>>(d_a);
  smem_6<<<1, 32>>>(d_a);
  cudaFree(d_a);
  cudaDeviceSynchronize();
  return 0;
}

总结

可以看到实际上 32 位宽的访存指令和 64/128 位宽的访存指令的广播机制以及 Bank Conflict 的计算都有很大的不同，但是官方文档中没有出现相关的描述（或者我没看到😵‍💫）。这篇文章通过 Microbenchmark 以及前人的一些讨论总结了几套规则，但需要注意的是这些规则有一定的局限性，其一是我只评测了以上图片中的例子，因此是不清楚更加复杂的访存情况是不是仍然符合这些规则的，其二是这些规则不是官方记录的，因此很有可能在将来被新发布的 GPU 架构所改写。不过能清楚这些底层细节还是很有趣的 ~~（大概～~~

参考

见 Volta GPU 的白皮书和 Turing GPU 的白皮书 ↩︎