[Hopper 架构特性学习笔记 Part2] Tensor Memory Access（TMA）

正文

请到「今天看啥」查看全文

是一个 constexpr 类型，该问题可以缓解。笔者尝试过将 num_per_threads 改为 constexpr int 类型，结果是循环会展开，LDG 指令个数为 num_per_threads 个，STS 指令会被合并（单个 STS 指令最大支持 128 bits 的数据拷贝）。

所以，针对以上问题，GPU 硬件上提升 Shared Memory 拷贝性能有三个思路：

1. 将 LDG 和 STS 指令合并为一种新指令 ，减少需要运行的指令数量；

2. 提供异步执行 Shared Memory 拷贝 ，减少循环间的等待；

3. 在循环中避免展开指令 ，仅一条指令完成拷贝计算，这样可以减少冗余操作。

Ampere 架构按照上述三种思路新增了一项新特性——异步拷贝。该特性支持异步从 Global Memory 拷贝到 Shared Memory，可以在拷贝期间执行其他计算操作（不能是读取Shared Memory相关的计算），这样实现拷贝和计算重叠，提升线程执行的性能。该特性还减少了寄存器“桥梁”的使用，无需再通过寄存器作为中间变量进行 Global Memory 到 Shared Memory 的拷贝。在如下图所示，在非异步拷贝的场景下，从 Global Memory 拷贝到 Shared Memory 需要经过 L2、L1、寄存器，而在异步拷贝场景下直接跳过寄存器。如果不需要 L1 这层存储器，在 PTX 层可以调用 cp.async.cg.shared.global.L2 绕开 L1，直接从 L2 拷贝到 Shared Memory 中。

下面将介绍异步读写的使用方式。

1.2 异步拷贝的使用

CUDA 提供了两种异步拷贝接口：cuda::memcpy_async 以及 cooperative_groups::memcpy_async。两种接口功能一样，下面以 cuda::memcpy_async 为例介绍异步拷贝的用法。

__global__ void AsyncCopyToSharedMem(int* data, int num_per_threads) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    auto grid = cooperative_groups::this_grid();
    auto block = cooperative_groups::this_thread_block();
    extern __shared__ int shm[];

    __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> barrier;
    if (block.thread_rank() == 0) {
      init(&barrier, block.size()); // Friend function initializes barrier
    }
    block.sync();
    #pragma unroll
    for (int i = 0; i < num_per_threads; ++i) {
        cuda::memcpy_async(block, shm + threadIdx.x * num_per_threads + i, data + idx * num_per_threads + i, sizeof(int), barrier);
    }
    barrier.arrive_and_wait(); // Waits for all copies to complete
}

代码中 cuda::memcpy_async 用以替代常规的 Shared Memory 赋值方式。调用时直接返回，不用等待拷贝完成，所以在循环中每次迭代都是立即返回，无需等待上次迭代的拷贝结束，这样解决了常规拷贝串行执行等待的问题。由于这个操作是异步的，常规的 __syncthreads() 无法感知异步拷贝是否完成，需要使用一种新的同步机制在合适的地方等待拷贝完成。代码中使用 mbarrier 完成异步拷贝的同步操作，在 cuda::memcpy_async 中设置 cuda::barrier 对象，并在后续的代码调用 barrier.arrive_and_wait() 等待当前线程的拷贝结果。cuda::barrier有一个计数器，可以统计完成异步拷贝的线程数，当 Thread Block 中所有线程完成拷贝后，则停止等待并返回。综上所述，异步拷贝的使用方式为：

1. 使用 cooperative_groups 获取当前 block；

2. 在 Shared Memory 上创建 cuda::barrier 对象，并使用一个线程对其初始化；

3. 调用 cuda::memcpy_async，将上一步创建的 cuda::barrier 对象放到 barrier 参数中；

4. 使用 barrier 同步异步拷贝。

在 cuda::memcpy_async 和 barrier.arrive_and_wait() 之间可以插入与 Shared Memory 无关的计算代码，从而实现计算与 Shared Memory 拷贝的重叠。

如上一节一样，我们通过 nvcc -o shm shm.cu -arch=compute_80 -code=sm_80 编译以上代码，并运行 cuobjdump --dump-sass shm 打印 SASS 层指令，得到以下代码：

/*02c0*/              @!PT LDS RZ, [RZ] ;                                    /* 0x00000000fffff984 */
                                                                             /* 0x000fe20000000800 */
/*02d0*/                   LDGSTS.E [R9], [R4.64] ;                          /* 0x0000000004097fae */
                                                                             /* 0x0003e2000b921844 */
/*02e0*/                   ISETP.GE.AND P0, PT, R8, 0x4, PT ;                /* 0x000000040800780c */
                                                                             /* 0x000fda0003f06270 */

通过代码可以看到，cuda::memcpy_async 会被编译成 LDGSTS 指令。这样实现了我们上面提到的优化思路：将 LDG 和 STS 指令合并为一种新指令，减少了指令数。并且该指令并没有展开成若干个指令，潜在可以降低冗余操作（这个结论要进一步探索，CUDA 官方文档提到了使用异步拷贝的优化点——对齐地址 16 字节。如果地址没对齐，硬件可能也会发起多次拷贝，性能可能会下降）。

小结 ：通过上面分析，Ampere 架构提供的异步拷贝可以降低指令数量（合并LDG 以及 STS指令），单个线程级别的 Shared Memory 拷贝和计算的重叠以及冗余操作。但是，针对一大块不连续的显存拷贝，需要开发者显式计算每段连续显存的首地址，从而引入了地址计算开销，所以每次执行 LDGSTS 前还是需要同步等待地址的计算。针对这个问题， Hopper 架构提出了 Tensor Memory Access（TMA），可以减少地址计算的开销。

2. Tensor Memory Access（TMA）

无论是常规的 Shared Memory 拷贝，还是 Ampere 架构下的 Shared Memory 异步拷贝，在拷贝大块的显存时，都会拆分成若干个很小的显存块，利用循环、多线程方式完成多个小显存块拷贝。每次拷贝均要计算显存的起始地址，这种寻址操作是不能被异步拷贝重叠的，并且运算指令随着小显存块的增多而线性增加。显式计算地址的原因主要是地址不连续，比如在矩阵乘中，对 Global Memory 进行分块，并将每个小块加载到 Shared Memory 中，显存块中不同行的地址是不连续的，需要手动计算。所以 Ampere 及其以前的架构，是无法减少这种频繁的地址计算操作。为了解决这个问题，Hopper 架构引入了 TMA 功能。TMA 支持以下几个功能：

1. 大块（bulk）异步显存拷贝 。使用 cuda::memcpy_async 接口。这个类似 CPU 上的 memcpy，支持一整块的显存拷贝，可以减少拷贝指令数量；

2. 多维度显存块拷贝 。 这个特性主要支持不连续的多段显存块拷贝。在实际使用中，需要区分一维度显存块拷贝和多维度显存块拷贝。多维度显存块拷贝需要在 Host 端调用 cuTensorMapEncode 的 API ，计算显存块之间的地址映射关系，然后通过带有__grid_constant__ 注释的 CUtensorMap 类型参数传递给 Kernel 函数中，调用 TMA 的异步拷贝接口完成多维度的拷贝。如下图所示，TMA 可以支持一个拷贝指令完成粉色显存块的拷贝；

请到「今天看啥」查看全文