原文：https://zhuanlan.zhihu.com/p/709750258
上一篇 [Hopper 架构特性学习笔记 Part1] 介绍了 Hooper 架构带来两个重点的新特性：Distributed Shared Memory 和 TMA ，并重点介绍了 Distributed Shared Memory 的使用。本篇继续介绍 Hooper 架构带来的第二个新特性：Tensor Memory Access (TMA)。 在介绍 TMA 之前，先介绍下 Ampere 架构带来的一个新特性——异步拷贝。
1. 异步拷贝 如上篇[Hopper 架构特性学习笔记 Part1] 提到，在 GPU 编程中，Kernel 的设计往往是以 Thread Block 这个粒度展开的。在 Thread Block 这个层次设计 Kernel 时，可以将重复使用的 Global Memory 拷贝到 Shared Memory，然后执行计算，计算结果在 Shared Memory上，再把结果写回到 Global Memory 上，利用 Shared Memory 访问延迟低的特点提升计算过程中的访存速度。所以在这种 Kernel 中，存在着频繁从 Global Memory 拷贝到 Shared Memory。如果可以优化 Global Memory 和 Shared Memory 之间的拷贝带来的开销，就可以优化 Kernel 的性能。我们通过一个简单的 Kernel 代码示例展示 Shared Memory 的使用方式，并打印其 SASS 指令，看看 GPU 在指令层级是如何实现的一个简单的 Shared Memory 拷贝。
1.1 同步拷贝存在的问题 // shm.cu void CopyToSharedMem (int * global_data) {int idx = threadIdx.x + blockIdx.x * blockDim.x;extern __shared__ int shm[];= global_data[idx];// ...... 这是一个最常见的 Shared Memory 的使用方式。在声明 Shared Memory 数组后，每个线程为 Shared Memory 数组中的一个元素赋值，然后调用 Thread Block 同步操作。同步调用结束后，再继续对 Shared Memory 进行计算操作。其中 shm[threadIdx.x]= global_data[idx] 这一行调用就是把 Global Memory 拷贝到 Shared Memory 中。在 GPU 中，这个拷贝操作会被分解为三个步骤：
计算 Global Memory 和 Shared Memory 的地址；
把 Global Memory 拷贝到寄存器中；
把寄存器的值拷贝到 Shared Memory 中。
通过 nvcc -o shm -arch=compute_80 -code=sm_80 cuobjdump --dump-sass shm
Function : _Z15CopyToSharedMemPiSASS 中有两个指令展示了 Global Memory 拷贝到 Shared Memory 过程，分别是LDG.E.CONSTANT R2, [R2.64] STS [R5.X4], R2 IMAD
// shm.cu void CopyToSharedMem (int * data, int num_per_threads) {int idx = threadIdx.x + blockIdx.x * blockDim.x;extern __shared__ int shm[];for (int i = 0 ; i < num_per_threads; ++ i) {* num_per_threads + i] = data[num_per_threads * idx + i];// ...... 其中 num_per_threads
/*0180*/ @!P1 BRA 0x410 ; /* 0x0000028000009947 */通过 SASS 指令可以观察到，在一次循环中 LDG 和 STS 指令各自展开成 16 个（应该是为了对齐地址），并且依次执行。综合上述两个 Kernel 示例可以发现，常规的 Shared Memory 赋值方式存在以下几个问题：
一次赋值涉及多个指令执行 。地址计算、LDG 以及 STS指令。LDG 和 STS 指令是通过一个寄存器作为 “桥梁”，实现 Shared Memory 赋值。这种指令设计我理解应该跟 RISC 架构有关。GPU 是 Load/Store 架构，每个指令的操作数必有一个是寄存器，所以需要使用一个寄存器作为“桥梁”；
单一线程中无数据依赖的 Load/Store 指令会串行同步执行。 单个线程内部执行拷贝 Global Memory 的循环迭代中，不同轮次的迭代并无数据依赖，同步执行会造成线程内频繁等待赋值指令完成的情况；
存在冗余的操作 。循环内部 Shared Memory 赋值会从原来 2 个指令展开为 32个指令，假如循环只有一次，也会执行 32 个指令，冗余执行了多个无用的指令，浪费时钟周期。假如循环次数 num_per_threads 是一个 constexpr 类型，该问题可以缓解。笔者尝试过将 num_per_threads 改为 constexpr int 类型，结果是循环会展开，LDG 指令个数为 num_per_threads 个，STS 指令会被合并（单个 STS 指令最大支持 128 bits 的数据拷贝）。
所以，针对以上问题，GPU 硬件上提升 Shared Memory 拷贝性能有三个思路：
将 LDG 和 STS 指令合并为一种新指令 ，减少需要运行的指令数量；
提供异步执行 Shared Memory 拷贝 ，减少循环间的等待；
在循环中避免展开指令 ，仅一条指令完成拷贝计算，这样可以减少冗余操作。
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[];:: barrier< cuda:: thread_scope:: thread_scope_block> barrier;if (block.thread_rank() == 0 ) {& barrier, block.size()); // Friend function initializes barrier }#pragma unroll for (int i = 0 ; i < num_per_threads; ++ i) {:: memcpy_async(block, shm + threadIdx.x * num_per_threads + i, data + idx * num_per_threads + i, sizeof (int ), barrier);// 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 中所有线程完成拷贝后，则停止等待并返回。综上所述，异步拷贝的使用方式为：
使用 cooperative_groups 获取当前 block；
在 Shared Memory 上创建 cuda::barrier 对象，并使用一个线程对其初始化；
调用 cuda::memcpy_async，将上一步创建的 cuda::barrier 对象放到 barrier 参数中；
使用 barrier 同步异步拷贝。
在 cuda::memcpy_async 和 barrier.arrive_and_wait() 之间可以插入与 Shared Memory 无关的计算代码，从而实现计算与 Shared Memory 拷贝的重叠。
如上一节一样，我们通过 nvcc -o shm -arch=compute_80 -code=sm_80 cuobjdump --dump-sass shm
/*02c0*/ @!PT LDS RZ, [RZ] ; /* 0x00000000fffff984 */通过代码可以看到，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 支持以下几个功能：
大块（bulk）异步显存拷贝 。使用 cuda::memcpy_async 接口。这个类似 CPU 上的 memcpy，支持一整块的显存拷贝，可以减少拷贝指令数量；
多维度显存块拷贝 。这个特性主要支持不连续的多段显存块拷贝。在实际使用中，需要区分一维度显存块拷贝和多维度显存块拷贝。多维度显存块拷贝需要在 Host 端调用 cuTensorMapEncode 的 API，计算显存块之间的地址映射关系，然后通过带有__grid_constant__ 注释的 CUtensorMap 类型参数传递给 Kernel 函数中，调用 TMA 的异步拷贝接口完成多维度的拷贝。如下图所示，TMA 可以支持一个拷贝指令完成粉色显存块的拷贝；
支持从 Shared Memory 异步拷贝到 Global Memory 。Ampere架构只支持从 Global Memory 异步拷贝到 Shared Memory，而在 Hopper 架构上更进一步支持反向的拷贝操作，提升 Kernel 的在不同存储结构上的读写性能。
TMA 针对一维度（连续地址）和多维度（不连续地址）显存块拷贝有着不同的用法，下面分别介绍。
2.1 一维度显存块拷贝 下面先通过一个代码示例介绍 TMA 的用法。
#include using barrier = cuda:: barrier< cuda:: thread_scope_block> ;static constexpr size_t buf_len = 1024 ;void add_one_kernel (int * data, size_t offset) {// Shared memory 数组。数组整体 size 要对齐 16字节 __shared__ alignas (16 ) int smem_data[buf_len];// 1. a) 用0号线程初始化 barrier，与上面的代码示例类似。 // b) 插入一个fence。表示后续执行异步拷贝操作，需要在这个fence之后才执行。 #pragma nv_diag_suppress static_var_with_dynamic_init __shared__ barrier bar;if (threadIdx.x == 0 ) { & bar, blockDim.x); // a) cuda:: device:: experimental:: fence_proxy_async_shared_cta();// b) }// 2. 发起 TMA 异步拷贝。注意：TMA 操作是用单线程发起。 if (threadIdx.x == 0 ) {// 3a. 发起异步拷贝 cuda:: memcpy_async(+ offset, :: aligned_size_t< 16 > (sizeof (smem_data)),// 3b. 所有线程到达该标记点，barrier内部的计数器会加 1。 barrier:: arrival_token token = bar.arrive();// 3c.等待barrier内部的计数器等于期望数值，即所有线程到达3b点时，当前线程的wait会返回，结束等待。 bar.wait(std:: move(token));// 4. 在 Shared Memory 上写数据。 for (int i = threadIdx.x; i < buf_len; i += blockDim.x) {+= 1 ;// 5. 插入fence，保证后续的异步拷贝操作在Shared Memory写数据结束后再启动。 cuda:: device:: experimental:: fence_proxy_async_shared_cta(); // b) __syncthreads();// 6. 发起从 Shared Memory 到 Global Memory 的异步拷贝操作。 if (threadIdx.x == 0 ) {:: device:: experimental:: cp_async_bulk_shared_to_global(+ offset, smem_data, sizeof (smem_data));// 7. 一种同步方式，创建一个 bulk async-group，异步拷贝在这个 group 中运行，当异步拷贝结束后， // group 内部标记为已完成。 cuda:: device:: experimental:: cp_async_bulk_commit_group();// 等待 group 完成。模版参数 0 表示要等待小于等于 0 个 bulk async-group 完成才结束等待。 cuda:: device:: experimental:: cp_async_bulk_wait_group_read< 0 > ();上述代码展示了 TMA 两种拷贝，分别是从 Global Memory 到 Shared Memory 的拷贝以及从 Shared Memory 到 Global Memory 的拷贝：
对于 Global Memory 到 Shared Memory 的拷贝，用法与 Ampere 架构下的异步拷贝类似，都使用 cuda::memcpy_async 接口发起异步拷贝操作，并使用 mbarrier 进行同步。但 TMA 的拷贝有两点不同：
1. 需要调用 fence_proxy_async_shared_cta 函数 。在异步拷贝前插入fence，保证必须在 fence 前某个操作完成才发起异步拷贝；
2. 单线程发起 。这其实对应了上面提到的大块异步显存拷贝的概念。整个 Thread Block 只发起了一个异步拷贝操作，减少拷贝的指令数。
对于 Shared Memory 到 Global Memory 的异步拷贝，则需要使用 cp_async_bulk 系列的接口了。这些接口都定义在 头文件里，是内联汇编的 Wrapper 函数。 cp_async_bulk_shared_to_global 发起了从 Shared Memory 到 Global Memory 的异步拷贝操作。不同于以往的同步方式， Shared Memory 到 Global Memory 的异步拷贝是通过 bulk async-group 方式同步的。在发起完异步拷贝操作后，需要用 cp_async_bulk_commit_group 创建一个 bulk async-group，确保异步拷贝在这个 group 中进行。当异步拷贝结束后，group 内部标记为已完成。最后，通过调用 cp_async_bulk_wait_group_read 完成同步操作。该函数是模版函数，模板参数为 group 数量，表示未完成异步拷贝的 group 数量小于等于这个 group 数量就会结束等待操作，如果 group 数量为 0，表示等待所有 group 完成才结束等待。
所以整个一维度的 TMA 使用流程，可以总结为以下三个阶段共计 6 个步骤：
--第一阶段： Global Memory 到 Shared Memory 的拷贝：
【初始化同步变量 】使用单线程初始化 Shared Memory 上的 barrier 对象；
【保证异步操作的顺序 】插入 fence，保证 barrier 对象初始化执行完才开始后续的异步拷贝；
【调用 cuda::memcpy_async&同步 】使用单线程发起异步拷贝，并通过 barrier.arrive 以及 barrier.wait 进行同步
--第二阶段： 对 Shared Memory 进行写操作。
--第三阶段：Shared Memory 到 Global Memory 的拷贝：
【保证异步操作的顺序】 插入 fence，保证前面的 Shared Memory 写操作完成后才开始后续的异步拷贝；
【 调用cp_async_bulk_shared_to_global & 同步 】使用单线程调用 cp_async_bulk_shared_to_global 发起异步拷贝，通过 cp_async_bulk_commit_group 以及 cp_async_bulk_wait_group_read 完成同步。
在了解完整个一维度的 TMA 使用流程后，我们进一步了解GPU底层具体做了什么。通过 nvcc -o shm -arch=compute_90 -code=sm_90 cuobjdump --dump-sass shm
// ......通过 SASS 层指令可以看到，TMA 的拷贝分别被编译为 UBLKCP.S.G 以及 UBLKCP.G.S 指令，对应不同方向的异步拷贝操作。由于整个 Thread Block 只发起了一个拷贝指令，所以指令的操作数放在 Uniform 寄存器中而非一般的线程私有的寄存器。Uniform 寄存器是 Warp 内线程共享的寄存器，Warp 内对相同的操作数进行操作时使用 Uniform 寄存器可以减少寄存器使用。对于一维度显存拷贝，地址计算仍然在线程中同步进行，没有重叠计算。不过一维度显存拷贝只有一次地址计算，重叠计算的收益也不会很大。
2.2 多维度显存块拷贝
上面提到，在多维度显存块拷贝中，需要在 Kernel 内部为每一小段连续显存计算其首地址，以保证拷贝正确的显存段。但是当维度很高，显存段的段数很多时，Kernel 内需要频繁计算地址，这部分计算开销不容忽视。实际当Global Memory 的维度、Shape以及Shared Memory 的 维度、Shape 确定时，Kernel 中每个 Thread Block 所要拷贝的 Global Memory 的首地址都能提前确定。CUDA 12 提供 CUtensorMap 结构，在 Host 端通过指定 Global Memory 的维度、Shape以及首地址、Shared Memory 的 维度、Shape 初始化 CUtensorMap 结构，即可提前计算出不同 Thread Block 对应的 Global Memory 的不同段的首地址。如上图右侧所示，地址不需要再实时计算。以下是 CUtensorMap 的初始化方式：
CUtensorMap tensor_map{};// rank is the number of dimensions of the array. constexpr uint32_t rank = 2 ;uint64_t size[rank] = {GMEM_WIDTH, GMEM_HEIGHT};// The stride is the number of bytes to traverse from the first element of one row to the next. // It must be a multiple of 16. uint64_t stride[rank - 1 ] = {GMEM_WIDTH * sizeof (int )};// The box_size is the size of the shared memory buffer that is used as the // destination of a TMA transfer. uint32_t box_size[rank] = {SMEM_WIDTH, SMEM_HEIGHT};// The distance between elements in units of sizeof(element). A stride of 2 // can be used to load only the real component of a complex-valued tensor, for instance. uint32_t elem_stride[rank] = {1 , 1 };// Get a function pointer to the cuTensorMapEncodeTiled driver API. auto cuTensorMapEncodeTiled = get_cuTensorMapEncodeTiled();// Create the tensor descriptor. CUresult res = cuTensorMapEncodeTiled(& tensor_map, // CUtensorMap *tensorMap, CUtensorMapDataType:: CU_TENSOR_MAP_DATA_TYPE_INT32,// cuuint32_t tensorRank, tensor_ptr, // void *globalAddress, size, // const cuuint64_t *globalDim, stride, // const cuuint64_t *globalStrides, box_size, // const cuuint32_t *boxDim, elem_stride, // const cuuint32_t *elementStrides, // Interleave patterns can be used to accelerate loading of values that // are less than 4 bytes long. CUtensorMapInterleave:: CU_TENSOR_MAP_INTERLEAVE_NONE,// Swizzling can be used to avoid shared memory bank conflicts. CUtensorMapSwizzle:: CU_TENSOR_MAP_SWIZZLE_NONE,// L2 Promotion can be used to widen the effect of a cache-policy to a wider // set of L2 cache lines. CUtensorMapL2promotion:: CU_TENSOR_MAP_L2_PROMOTION_NONE,// Any element that is outside of bounds will be set to zero by the TMA transfer. CUtensorMapFloatOOBfill:: CU_TENSOR_MAP_FLOAT_OOB_FILL_NONECUtensorMap 对象通过 cuTensorMapEncodeTiled 初始化，需要指定 Global Memory 首地址（tensor_ptr）、Shape(size)、 Shared Memory 的 Shape（box_size）以及维度（rank）。在完成 CUtensorMap 创建后，可以将其传入到 Kernel 内部使用，这样 Kernel 内部无需再计算地址了。以下是 TMA 多维度显存块拷贝的示例代码：
#include // CUtensormap#include using barrier = cuda:: barrier< cuda:: thread_scope_block> ;namespace cde = cuda:: device:: experimental;void kernel (const __grid_constant__ CUtensorMap tensor_map, int x, int y) {// bluk tensor 的拷贝操作需要 Shared Memory 首地址对齐 128 字节。 __shared__ alignas (128 ) int smem_buffer[SMEM_HEIGHT][SMEM_WIDTH];// 创建 Shared Memory 的 cuda::barrier 变量 #pragma nv_diag_suppress static_var_with_dynamic_init __shared__ barrier bar;if (threadIdx.x == 0 ) {// 初始化 barrier init(& bar, blockDim.x);// 插入 fence cde:: fence_proxy_async_shared_cta(); :: arrival_token token;if (threadIdx.x == 0 ) {// 发起 TMA 二维异步拷贝操作 cde:: cp_async_bulk_tensor_2d_global_to_shared(& smem_buffer, & tensor_map, x, y, bar);// 设置同步等待点，指定需要等待的拷贝完成的字节数。 token = cuda:: device:: barrier_arrive_tx(bar, 1 , sizeof (smem_buffer));else {// Other threads just arrive. token = bar.arrive();// 等待完成拷贝 bar.wait(std:: move(token));0 ][threadIdx.x] += threadIdx.x;// 插入 fence cde:: fence_proxy_async_shared_cta();if (threadIdx.x == 0 ) {:: cp_async_bulk_tensor_2d_shared_to_global(& tensor_map, x, y, & smem_buffer);:: cp_async_bulk_commit_group();:: cp_async_bulk_wait_group_read< 0 > ();if (threadIdx.x == 0 ) {& bar)->~ barrier(); 多维度显存块拷贝的使用范式与一维度显存块拷贝的使用范式基本一样，只有一些具体调用的函数不同，这里讲讲不同点：
不再使用cuda::memcpy_async 实现 Global Memory 到 Shared Memory 的拷贝，取而代之使用 cp_async_bulk_tensor_2d_global_to_shared 调用实现。该调用传入了 tensor_map 参数，根据 tensor_map 可以获取需要拷贝的 Global Memory 段的首地址；
同步方式。主要是调用异步拷贝（Global Memory 到 Shared Memory）的线程需要使用cuda::device::barrier_arrive_tx 调用，该调用可以传入拷贝的字节数，表示 barrier 需要等待多少字节拷贝完毕才结束。
如上一节一样，我们通过 nvcc -o shm -arch=compute_90 -code=sm_90 编译以上代码，并运行 cuobjdump --dump-sass shm 打印 SASS 层指令，得到以下代码：
// ......通过 SASS 层指令可以看到，TMA 的拷贝分别被编译为 UTMALDG.2D 以及 UTMASTG.2D 指令，分别对应 TMA LOAD （Global Memory To Shared Memory）和 TMA READ（Shared Memory To Global Memory） 操作。操作数地址可以通过读取 tensor_map 的值获取，避免了实时计算操作。以上代码展示二维的异步拷贝。Hopper 架构一共支持 5 个维度的异步拷贝，相关的函数如下：
3. 总结 本文主要介绍 Hopper 架构的一个新特性——Tensor Memory Access（TMA），并介绍在 CUDA 上如何使用 TMA。按 CUDA C++ Best Practices Guide （https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#asynchronous-copy-from-global-memory-to-shared-memory）的说法，异步拷贝基本在所有场景下都能比同步拷贝表现出更高的性能，而 TMA 是异步拷贝的加强版，由TMA单元计算地址，无需再 Kernel 内部计算地址，实现更高效的异步拷贝。CUTLASS 针对 Hopper 架构已经引入 TMA Tensor（https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/0z_tma_tensors.md），用于加速矩阵乘的读写 Shared Memory 效率，Flash Attention3 也使用了 CUTLASS 提供的 TMA 接口，相信 TMA 在以后的 Kernel 优化工作中出镜率会越来越高。
4. 参考 https:// research.colfax-intl.com/tutorial-hopper-tma/
https:// arxiv.org/html/2402.13499v1
https:// docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tensor-memory-access
PTX ISA 8.5
globalToShmemAsyncCopy.cu sample
CUDA C++ Best Practices Guide
reed：NVidia GPU指令集架构-寄存器
https:// docs.nvidia.com/cuda/cu
