主要观点总结
本文介绍了如何使用NVIDIA Hopper架构中的张量内存加速器(TMA)进行GPU内存操作，包括TMA load、TMA store、TMA store reduce和TMA load multicast。文章通过简化示例展示了如何在CUDA kernel中调用这些操作，并深入研究了底层PTX指令，以获得对TMA更深入的理解。文章还讨论了TMA的某些高级特性，如swizzling模式和以interleaved格式排列的能力，这些在使用TMA配合Warpgroup矩阵-乘法-累加(WGMMA)指令时很重要。文章最后提供了完整的kernel示例。
关键观点总结
关键观点1: TMA简介与优势
TMA是NVIDIA Hopper架构中引入的一项新功能，用于在GPU的全局内存(GMEM)和其线程块(即CTA)的共享内存(SMEM)之间进行异步内存复制。TMA提供了许多优势，例如提高GPU利用率和节省寄存器。
关键观点2: TMA操作在CUDA kernel中的使用
文章通过简化示例展示了如何在CUDA kernel中调用TMA load、TMA store、TMA store reduce和TMA load multicast操作。
关键观点3: TMA操作底层的PTX指令
文章深入研究了底层PTX指令，以获得对TMA更深入的理解，并讨论了TMA支持的swizzling模式以及以interleaved格式排列的能力。
关键观点4: 完整的kernel示例
文章最后提供了完整的kernel示例，这些示例涵盖了文章中讨论的所有TMA操作。
正文
它的第一个参数是SM90_TMA_LOAD(https://github.com/NVIDIA/cutlass/blob/637b15906358191cb4238af419d408a65819d7ec/include/cute/arch/copy_sm90_tma.hpp#L269)的一个实例。这个对象将复制操作分派到所需的
cp.async.bulk.tensor
PTX调用，我们将在下面的第三部分中深入探讨。
Kernel code
相关的 kernel 代码片段如下所示。这些代码行包含了许多重要的TMA概念，我们将在下面进行解释。
首先，在第2行，kernel的tma_load参数必须用
__grid_constant__ const
注解。如果我们有两个要从GMEM复制到SMEM的张量，每个张量都必须有自己的
TiledCopy
实例，并且每个实例都必须是
__grid_constant__ const
。这是从主机传递
cuTensorMap
到设备的要求，例如在这里有文档(https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copies-using-tensor-memory-access-tma)说明。
下一个重要点是，对于TMA Copy，只有一个线程负责发出TMA操作。在代码片段中，所有与TMA相关的变量和指令都包含在从第12行开始的if块中，该块仅由线程0执行。另一方面，第30行包含一条指令，让CTA中的所有线程等待TMA操作完成。
坐标和算术元组
现在，让我们看看TMA load逻辑。这从第13行开始，我们创建一个
gmem_tensor_coord
对象，它保存要复制的GMEM张量的坐标。如果我们尝试以下操作：
if (cute::thread(0)) { cute::print(gmem_tensor_coord); }
那么我们会看到如下输出（对于M=N=1024）：
ArithTuple(_0,_0) o (1024,1024):(_1@1,_1@0)
对于熟悉CuTe中tiled copy工作方式的读者来说，第15-18行是不言自明的，其中GMEM张量被tiled成更小的partitions，每个CTA根据块坐标切片到tiled张量中以获得其GMEM视图。但是请注意，partitions适用于上述表示gmem_tensor坐标的ArithTuple，而不是gmem_tensor本身。特别是，ArithTuple被分成形状为
[CTA_M,CTA_N]
的块，然后每个CTA取其块。
如果我们使用
print_tensor
打印
gmem_tensor_coord_cta
，如下所示：
if (cute::block(7)) { cute::print_tensor(gmem_tensor_coord_cta); }
我们会看到如下输出：
ArithTuple(0,112) o (_16,_16):(_1@1,_1@0):
(0,112) (1,112) (2,112) (3,112) (4,112) (5,112) (6,112) (7,112) (8,112) (9,112) (10,112) (11,112) (12,112) (13,112) (14,112) (15,112)
(0,113) (1,113) (2,113) (3,113) (4,113) (5,113) (6,113) (7,113) (8,113) (9,113) (10,113) (11,113) (12,113) (13,113) (14,113) (15,113)
// more lines
(0,127) (1,127) (2,127) (3,127) (4,127) (5,127) (6,127) (7,127) (8,127) (9,127) (10,127) (11,127) (12,127) (13,127) (14,127) (15,127)
这些数字是
gmem_tensor
中的坐标，其值将被复制到CTA 7的
smem_tensor
中。我们鼓励读者尝试运行这段代码片段，将
cute::block(7)
替换为其他索引，以理解不同的CTA从
gmem_tensor
的哪些坐标复制数据。
接下来，在第25-27行发出的复制操作本身具有TiledCopy操作的常见签名，其中源张量被partitions后的坐标所替代。
Memory barrier
我们省略了第20、22和30行，这些行都涉及SMEM中的
uint64_t
变量
tma_load_mbar
。这是我们用来同步TMA load 与 kernel 消费 load 到SMEM中的结果数据的其余部分的
异步事务屏障
。NVIDIA关于Hopper架构的技术博客(https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/)中给出了这种屏障的高级描述。就我们的kernel而言，重要的点如下：
-
我们在第20行的共享内存中初始化mbarrier对象。CuTe方法
initialize_barrier
包装了PTX指令
mbarrier.init.shared.b64
，该指令需要一个额外的到达计数参数。在我们的上下文中，由于单个线程将启动TMA load，我们应该将到达计数设置为1。此外，mbarrier的起始阶段将始终设置为0。
-
我们在第22行同时执行arrive-on操作并为mbarrier对象设置预期的事务计数，使用CuTe方法
set_barrier_transaction_bytes
，它包装了PTX指令
mbarrier.arrive_expect_tx.shared::cta.b64
。事务计数设置为等于TMA load传输的字节数，我们在第4行计算这个值。
-
在第25-27行，复制指令（它分派到所需的
cp.async.bulk.tensor
类型）总是将其完成机制设置为
barrier::complete_tx::bytes
，并使用提供的mbarrier对象。
-
在第30行，我们在mbarrier对象上执行等待操作。注意，所有线程都在mbarrier上等待，这与只有线程0到达mbarrier形成对比，并且在
wait_barrier
之前调用
__syncthreads()
是必要的，以解决线程分歧。这里，
wait_barrier
包装了PTX指令
mbarrier.try_wait.parity.shared::cta.b64
。
try_wait
限定符（与
test_wait
相对）表示等待是一个阻塞指令。
parity
限定符（其使用需要提供一个相位）表示线程睡眠直到mbarrier的那个相位翻转。因为这是初始化后首次使用mbarrier来跟踪完成，我们提供0作为相位。如果我们要进行另一次TMA load，我们就需要翻转相位以重用mbarrier。总的来说，CUTLASS Pipeline APIs(https://github.com/NVIDIA/cutlass/blob/main/media/docs/pipeline.md)提供了一种更高级的方式来处理一系列TMA load时mbarrier对象的生命周期，就像在软件流水线(https://github.com/NVIDIA/cutlass/blob/main/media/docs/efficient_gemm.md#pipelining)方案中可能做的那样。
