1 PTX (Parallel Thread Execution)

PTX是什么，Nvidia官方描述为a low-level parallel thread execution virtual machine and instruction set architecture (ISA)，直面意思是低级并行线程执行虚拟机和指令集架构。

怎么理解其直面意思，有两个方法。

一个方法是借鉴LLVM，熟悉LLVM的知道其全称是Low Level Virtual Machine，这里不在意LLVM的主干项目与其底层虚拟机的命名渐行渐远，主要关注LLVM的核心概念IR（Intermediate Representation），其行为与PTX有几分相似。IR连接了前端编程语言和后端目标代码，不仅可以比较容易地实现新的编程语言，还可以方便地生成不同硬件平台上的目标代码，同时还可以做一些通用性的编译优化和运行时优化。PTX是上承GPU编程语言CUDA C++，下启GPU硬件SASS指令，可以借助NVRTC实现运行时优化，某些层面上来说可以称之为GPU设备无关代码，因此PTX可以理解为”CUDA IR“。

另一个方法是不用太理解，毕竟Nvidia闭源的出发点就是让开发者难得糊涂。

再回到PTX本身，习惯了CUDA C++编程，PTX似乎不曾看到过，但它其实一直都在。如下图所示为NVCC编译CUDA的过程，可以发现.cu文件的编译分为两个部分，一部分是编译主机代码，另一部分是编译设备代码，设备代码的编程过程中会生成.ptx文件，而通常关注的是编译生成的最终产物。NVCC的编译流程在这里就不展开了，后续有机会再聊。

2 MMA (Matrix Multiply Accumulate) PTX

对于计算能力在7.0及以上的CUDA设备，可以使用MMA PTX指令调用Tensor Core，支持形如D = AB + C的混合精度的矩阵乘运算。

mma.sync.aligned.m8n8k4.alayout.blayout.dtype.f16.f16.ctype  d, a, b, c;
mma.sync.aligned.m16n8k8.row.col.dtype.f16.f16.ctype  d, a, b, c;
mma.sync.aligned.m16n8k16.row.col.dtype.f16.f16.ctype d, a, b, c;

以m16n8k16 FP16为例，每个tile中的元素在warp内线程上的计算分布如下图所示，可以明显发现每个线程计算的fragment都是不连续的。

矩阵A fragment的行和列的索引可以按如下方式计算:

groupID           = %laneid >> 2
threadID_in_group = %laneid % 4row = groupID      for ai where 0 <= i < 2 || 4 <= i < 6groupID + 8  Otherwisecol = (threadID_in_group * 2) + (i & 0x1)      for ai where i < 4(threadID_in_group * 2) + (i & 0x1) + 8  for ai where i >= 4

矩阵B fragment的行和列的索引可以按如下方式计算:

groupID           = %laneid >> 2
threadID_in_group = %laneid % 4row = (threadID_in_group * 2) + (i & 0x1)      for bi where i < 2 (threadID_in_group * 2) + (i & 0x1) + 8  for bi where i >= 2col = groupID

矩阵C或D fragment的行和列的索引可以按如下方式计算:

groupID           = %laneid >> 2
threadID_in_group = %laneid % 4row = groupID      for ci where i < 2groupID + 8  for ci where i >= 2col = (threadID_in_group * 2) + (i & 0x1)  for ci where i = {0,..,3}

3 LDMATRIX PTX

由于MMA PTX指令计算tile时，warp内线程计算的fragment不连续，索引计算较为复杂，所以Nvidia提供了LDMATRIX PTX指令用来配合MMA PTX指令。

ldmatrix.sync.aligned.shape.num{.trans}{.ss}.type r, [p];.shape  = {.m8n8};
.num    = {.x1, .x2, .x4};
.ss     = {.shared};
.type   = {.b16};

LDMATRIX PTX是warp级别的数据加载指令，其读取连续的行不需要连续地存储在内存中。每个矩阵所需的8个地址由8个线程提供，具体取决于.num的值。每个地址对应于一个矩阵行的开始。地址addr0-addr7对应第一个矩阵的行，地址addr8-addr15对应第二个矩阵的行，依此类推，如下表所示。

当读取8x8的矩阵时，一组连续的四个线程加载16个字节。矩阵地址必须相应地对齐。warp中的每个线程加载一行的fragment，线程0接收寄存器r中的第一个fragment，以此类推。由四个线程组成的一组将加载矩阵的一整行，如下表所示。可以发现，LDMATRIX PTX指令在warp内线程上的数据分布与MMA PTX指令一致。

值得注意的是，首先LDMATRIX PTX指令只能从shared memory中加载数据；其次对于计算能力在sm_75及以下的CUDA设备，LDMATRIX PTX指令中的所有线程必须包含有效地址。否则，行为是未定义的。.num为.x1和.x2时，低线程中包含的地址可以复制到高线程中，以实现预期的行为。

4 示例

Talk is cheap，show me the code。与Nvidia Tensor Core-WMMA API编程入门类似，以m16n8k16为例，实现HGEMM：C = AB，其中矩阵A（M * K，row major）、B（K * N，col major）和C（M * N，row major）的精度均为FP16。

MMA PTX的编程思路类似于WMMA API，都是按照每个warp处理一个矩阵C的tile的思路来构建naive kernel。首先确定当前warp处理矩阵C的tile坐标，声明计算tilie所需的shared memory和寄存器，再以MMA_K为步长遍历K并从global memory经shared memory由LDMATRIX PTX加载所需A、B矩阵tile到寄存器参与计算，最后将计算结果从寄存器经shared memory写回矩阵C。所有block计算完成之后即可得到矩阵C。这个例子有难度，但不多。

#include <mma.h>#define DIV_CEIL(x, y) (((x) + (y) - 1) / (y))#define MMA_M 16
#define MMA_N 8
#define MMA_K 16#define WARP_SIZE 32__global__ void mma16816NaiveKernel(const half *__restrict__ A, const half *__restrict__ B, half *__restrict__ C,size_t M, size_t N, size_t K) {const size_t K_tiles = DIV_CEIL(K, MMA_K);const size_t warpRow = blockIdx.x * MMA_M;const size_t warpCol = blockIdx.y * MMA_N;if (warpRow >= M || warpCol >= N) {return;}__shared__ half shmem_A[MMA_M][MMA_K];__shared__ half shmem_B[MMA_N][MMA_K];__shared__ half shmem_C[MMA_M][MMA_N];const size_t laneId = threadIdx.x % WARP_SIZE;uint32_t RA[4];uint32_t RB[2];uint32_t RC[2] = {0, 0};for (size_t i = 0; i < K_tiles; ++i) {*((int4 *)(&shmem_A[laneId / 2][0]) + laneId % 2) =*((int4 *)(&A[(warpRow + laneId / 2) * K + i * MMA_K]) + laneId % 2);if (laneId < MMA_N * 2) {*((int4 *)(&shmem_B[laneId / 2][0]) + laneId % 2) =*((int4 *)(&B[i * MMA_K + (warpCol + laneId / 2) * K]) + laneId % 2);}__syncthreads();uint32_t shmem_A_lane_addr = __cvta_generic_to_shared(&shmem_A[laneId % 16][(laneId / 16) * 8]);asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n": "=r"(RA[0]), "=r"(RA[1]), "=r"(RA[2]), "=r"(RA[3]): "r"(shmem_A_lane_addr));uint32_t shmem_B_lane_addr = __cvta_generic_to_shared(&shmem_B[laneId % 8][((laneId / 8) % 2) * 8]);asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n": "=r"(RB[0]), "=r"(RB[1]): "r"(shmem_B_lane_addr));asm volatile("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n": "=r"(RC[0]), "=r"(RC[1]): "r"(RA[0]), "r"(RA[1]), "r"(RA[2]), "r"(RA[3]), "r"(RB[0]), "r"(RB[1]), "r"(RC[0]), "r"(RC[1]));*((uint32_t *)(&shmem_C[laneId / 4][0]) + laneId % 4)     = RC[0];*((uint32_t *)(&shmem_C[laneId / 4 + 8][0]) + laneId % 4) = RC[1];__syncthreads();}if (laneId < MMA_M) {*((int4 *)(&C[(warpRow + laneId) * N + warpCol])) = *((int4 *)(&shmem_C[laneId][0]));}__syncthreads();
}void hgemmMmaNaive(half *A, half *B, half *C, size_t M, size_t N, size_t K) {dim3 block(WARP_SIZE);dim3 grid(DIV_CEIL(M, MMA_M), DIV_CEIL(N, MMA_N));mma16816NaiveKernel<<<grid, block>>>(A, B, C, M, N, K);
}

5 底层代码

我们再对上述MMA naive kernel做进一步探索，看一下它在RTX A6000（sm_86，CUDA 11.3）上对应的SASS。

      IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] S2R R36, SR_CTAID.Y IMAD.MOV.U32 R7, RZ, RZ, c[0x0][0x188] S2R R34, SR_CTAID.X IMAD.SHL.U32 R36, R36, 0x8, RZ IMAD.SHL.U32 R34, R34, 0x10, RZ ISETP.GE.U32.AND P0, PT, R36, c[0x0][0x180], PT ISETP.GE.U32.AND P1, PT, R34, c[0x0][0x178], PT ISETP.GE.U32.AND.EX P0, PT, RZ, c[0x0][0x184], PT, P0 ISETP.GE.U32.OR.EX P0, PT, RZ, c[0x0][0x17c], P0, P1
@P0   EXIT IADD3 R7, P0, R7, 0xf, RZ S2R R6, SR_TID.X ULDC.64 UR6, c[0x0][0x118] IMAD.X R8, RZ, RZ, c[0x0][0x18c], P0 SHF.R.U64 R7, R7, 0x4, R8 SHF.R.U32.HI R8, RZ, 0x4, R8 ISETP.NE.U32.AND P0, PT, R7, RZ, PT ISETP.NE.AND.EX P0, PT, R8, RZ, PT, P0 LOP3.LUT R33, R6, 0x1f, RZ, 0xc0, !PT
@!P0  BRA 0x7fc374fcb660 IMAD.SHL.U32 R0, R6, 0x20, RZ IADD3 R3, P0, R7, -0x1, RZ IMAD.SHL.U32 R5, R6, 0x4, RZ LOP3.LUT R2, R33, 0x10, RZ, 0xc0, !PT IMAD.SHL.U32 R4, R33, 0x4, RZ LOP3.LUT R31, R0, 0x1e0, RZ, 0xc0, !PT UMOV UR4, URZ LOP3.LUT R30, R0, 0xe0, RZ, 0xc0, !PT UMOV UR5, URZ ISETP.GE.U32.AND P1, PT, R3, 0x3, PT IMAD.SHL.U32 R3, R6, 0x2, RZ IADD3.X R0, R8, -0x1, RZ, P0, !PT IMAD.IADD R31, R31, 0x1, R2 LOP3.LUT R5, R5, 0xc, RZ, 0xc0, !PT IMAD.SHL.U32 R2, R6, 0x10, RZ LOP3.LUT R4, R4, 0x70, RZ, 0xc0, !PT CS2R R20, SRZ ISETP.GE.U32.AND.EX P1, PT, R0, RZ, PT, P1 LOP3.LUT R3, R3, 0x10, RZ, 0xc0, !PT IMAD.IADD R27, R4, 0x1, R5 SHF.R.U64 R5, R33, 0x1, RZ LOP3.LUT R0, R7, 0x3, RZ, 0xc0, !PT IMAD.IADD R30, R30, 0x1, R3 LOP3.LUT R32, R2, 0x10, RZ, 0xc0, !PT IADD3 R3, P2, R34, R5, RZ IADD3 R4, P3, R36, R5, RZ IMAD R32, R5, 0x20, R32 ISETP.NE.U32.AND P0, PT, R0, RZ, PT IMAD.X R5, RZ, RZ, RZ, P2 LOP3.LUT R2, R33, 0x1, RZ, 0xc0, !PT IMAD.X R6, RZ, RZ, RZ, P3 ISETP.NE.AND.EX P0, PT, RZ, RZ, PT, P0
@!P1  BRA 0x7fc374fcb390 IADD3 R24, P1, -R0, R7, RZ IMAD.SHL.U32 R26, R2, 0x10, RZ ISETP.GT.U32.AND P2, PT, R33, 0xf, PT CS2R R20, SRZ SHF.L.U64.HI R25, R2, 0x4, RZ UMOV UR4, URZ IADD3.X R7, R8, -0x1, RZ, P1, !PT UMOV UR5, URZ USHF.L.U32 UR8, UR4, 0x4, URZ IMAD R8, R5, c[0x0][0x188], RZ USHF.L.U64.HI UR9, UR4, 0x4, UR5 IMAD R11, R3, c[0x0][0x18c], R8 IMAD.U32 R12, RZ, RZ, UR8 IMAD.U32 R13, RZ, RZ, UR9 IMAD.WIDE.U32 R8, R3, c[0x0][0x188], R12 IMAD.IADD R9, R9, 0x1, R11 LEA R41, P1, R8, c[0x0][0x160], 0x1 LEA.HI.X R8, R8, c[0x0][0x164], R9, 0x1, P1 IADD3 R40, P1, R26, R41, RZ IMAD.X R41, R25, 0x1, R8, P1 LDG.E.128.CONSTANT R8, [R40.64] IMAD R15, R6, c[0x0][0x188], RZ IMAD.WIDE.U32 R12, R4, c[0x0][0x188], R12 IMAD R15, R4, c[0x0][0x18c], R15 LEA R39, P1, R12, c[0x0][0x168], 0x1 IMAD.IADD R13, R13, 0x1, R15 LEA.HI.X R12, R12, c[0x0][0x16c], R13, 0x1, P1 IADD3 R38, P1, R39, R26, RZ IMAD.X R39, R12, 0x1, R25, P1
@!P2  LDG.E.128.CONSTANT R16, [R38.64]
@!P2  LDG.E.128.CONSTANT R12, [R38.64+0x20] STS.128 [R32], R8 LDG.E.128.CONSTANT R8, [R40.64+0x20]
@!P2  STS.128 [R32+0x200], R16 BAR.SYNC 0x0 LDSM.16.M88.2 R22, [R30+0x200] LDSM.16.M88.4 R16, [R31] HMMA.16816.F16 R20, R16, R22, R20 NOP STS [R27+0x300], R20 STS [R27+0x380], R21 BAR.SYNC 0x0
@!P2  STS.128 [R32+0x200], R12
@!P2  LDG.E.128.CONSTANT R12, [R38.64+0x40] STS.128 [R32], R8 BAR.SYNC 0x0 LDSM.16.M88.2 R28, [R30+0x200] LDSM.16.M88.4 R16, [R31] LDG.E.128.CONSTANT R8, [R40.64+0x40] HMMA.16816.F16 R28, R16, R28, R20 LDG.E.128.CONSTANT R16, [R40.64+0x60]
@!P2  LDG.E.128.CONSTANT R20, [R38.64+0x60] NOP STS [R27+0x300], R28 STS [R27+0x380], R29 BAR.SYNC 0x0
@!P2  STS.128 [R32+0x200], R12 IADD3 R24, P1, R24, -0x4, RZ IADD3.X R7, R7, -0x1, RZ, P1, !PT ISETP.NE.U32.AND P1, PT, R24, RZ, PT ISETP.NE.AND.EX P1, PT, R7, RZ, PT, P1 UIADD3 UR4, UP0, UR4, 0x4, URZ UIADD3.X UR5, URZ, UR5, URZ, UP0, !UPT STS.128 [R32], R8 BAR.SYNC 0x0 LDSM.16.M88.2 R12, [R30+0x200] LDSM.16.M88.4 R8, [R31] HMMA.16816.F16 R12, R8, R12, R28 NOP STS [R27+0x300], R12 STS [R27+0x380], R13 BAR.SYNC 0x0 STS.128 [R32], R16
@!P2  STS.128 [R32+0x200], R20 BAR.SYNC 0x0 LDSM.16.M88.2 R14, [R30+0x200] LDSM.16.M88.4 R8, [R31] HMMA.16816.F16 R20, R8, R14, R12 NOP STS [R27+0x300], R20 STS [R27+0x380], R21 BAR.SYNC 0x0
@P1   BRA 0x7fc374fcaee0
@!P0  BRA 0x7fc374fcb660 SHF.L.U64.HI R11, R2, 0x3, RZ IMAD.SHL.U32 R10, R2, 0x8, RZ USHF.L.U32 UR8, UR4, 0x4, URZ IMAD R13, R6, c[0x0][0x188], RZ USHF.L.U64.HI UR4, UR4, 0x4, UR5 IMAD.WIDE.U32 R8, R4, c[0x0][0x188], R10 IADD3 R0, P3, RZ, -R0, RZ IMAD R2, R5, c[0x0][0x188], RZ IADD3 R22, P1, R8, UR8, RZ IMAD.WIDE.U32 R6, R3, c[0x0][0x188], R10 IMAD R13, R4, c[0x0][0x18c], R13 IMAD R5, R3, c[0x0][0x18c], R2 IADD3 R3, P0, R6, UR8, RZ IMAD.X R18, RZ, RZ, -0x1, P3 IADD3.X R9, R13, UR4, R9, P1, !PT LEA R19, P1, R22, c[0x0][0x168], 0x1 LEA R2, P2, R3, c[0x0][0x160], 0x1 IADD3.X R6, R5, UR4, R7, P0, !PT LEA.HI.X R22, R22, c[0x0][0x16c], R9, 0x1, P1 LEA.HI.X R3, R3, c[0x0][0x164], R6, 0x1, P2 ISETP.GT.U32.AND P0, PT, R33, 0xf, PT LDG.E.128.CONSTANT R4, [R2.64]
@!P0  IMAD.MOV.U32 R8, RZ, RZ, R19
@!P0  IMAD.MOV.U32 R9, RZ, RZ, R22
@!P0  LDG.E.128.CONSTANT R8, [R8.64] IADD3 R2, P2, R2, 0x20, RZ IADD3 R19, P1, R19, 0x20, RZ IMAD.X R3, RZ, RZ, R3, P2 IMAD.X R22, RZ, RZ, R22, P1 STS.128 [R32], R4
@!P0  STS.128 [R32+0x200], R8 BAR.SYNC 0x0 LDSM.16.M88.2 R16, [R30+0x200] LDSM.16.M88.4 R12, [R31] IADD3 R0, P0, R0, 0x1, RZ IMAD.X R18, RZ, RZ, R18, P0 ISETP.NE.U32.AND P0, PT, R0, RZ, PT ISETP.NE.AND.EX P0, PT, R18, RZ, PT, P0 HMMA.16816.F16 R20, R12, R16, R20 NOP STS [R27+0x300], R20 STS [R27+0x380], R21 BAR.SYNC 0x0
@P0   BRA 0x7fc374fcb4e0 ISETP.GT.U32.AND P0, PT, R33, 0xf, PT BSSY B0, 0x7fc374fcb760
@P0   BRA 0x7fc374fcb750 LEA R4, R33, 0x300, 0x4 IMAD.MOV.U32 R37, RZ, RZ, RZ IADD3 R33, P0, R34, R33, RZ LDS.128 R4, [R4] IMAD.X R0, RZ, RZ, RZ, P0 IMAD.WIDE.U32 R36, R33, c[0x0][0x180], R36 IMAD R0, R0, c[0x0][0x180], RZ LEA R2, P0, R36, c[0x0][0x170], 0x1 IMAD R3, R33, c[0x0][0x184], R0 IMAD.IADD R3, R37, 0x1, R3 LEA.HI.X R3, R36, c[0x0][0x174], R3, 0x1, P0 STG.E.128 [R2.64], R4 BSYNC B0 BAR.SYNC 0x0 EXIT BRA 0x7fc374fcb780NOPNOPNOPNOPNOPNOPNOPNOPNOPNOPNOPNOPNOPNOPNOP

可以发现与WMMA161616 API类似，MMA16816 PTX指令底层实现也是HMMA16816指令。

6 其他

6.1 HGEMM优化

与WMMA API类似，学习MMA PTX的目标在于调用Tensor Core优化HGEMM，相比于cublas，MMA的性能究竟如何？我们后续再聊。

Nvidia Tensor Core-MMA PTX编程入门相关推荐

深入理解混合精度训练：从 Tensor Core 到 CUDA 编程
背景近年来,自动混合精度(Auto Mixed-Precision,AMP)技术在各大深度学习训练框架中作为一种使用简单.代价低廉.效果显著的训练加速手段,被越来越广泛地应用到算法研究中. 然而大部 ...
Nvidia Tensor Core-WMMA API编程入门
1 WMMA (Warp-level Matrix Multiply Accumulate) API 对于计算能力在7.0及以上的CUDA设备,可以使用CUDA C++ API调用Tensor Cor ...
NVIDIA Tensor Cores解析
NVIDIA Tensor Cores解析高性能计算机和人工智能前所未有的加速 Tensor Cores支持混合精度计算,动态调整计算以加快吞吐量,同时保持精度.最新一代将这些加速功能扩展到各种工作 ...
NVIDIA深度学习Tensor Core性能解析（下）
NVIDIA深度学习Tensor Core性能解析(下) DeepBench推理测试之RNN和Sparse GEMM DeepBench的最后一项推理测试是RNN和Sparse GEMM,虽然测试中可 ...
NVIDIA深度学习Tensor Core性能解析（上）
NVIDIA深度学习Tensor Core性能解析(上) 本篇将通过多项测试来考验Volta架构,利用各种深度学习框架来了解Tensor Core的性能. 很多时候,深度学习这样的新领域会让人难以理解 ...
如何利用 NVIDIA 安培架构 GPU 的新一代 Tensor Core 对计算进行极致加速
时隔三年,英伟达最强芯片 Tesla V100 终于有了继任者,那就是在2020年5月14日, NVIDIA发布的最新GPU A100. GPU A100不仅能实现 1-50 倍的扩展,还让其吞吐性能 ...
core组件进阶（Opencv3编程入门第五章）第四节图像对比度亮度调整
core组件进阶(Opencv3编程入门第五章) 第四节图像对比度亮度调整 #include <opencv2/core/core.hpp> #include <opencv2/ ...
Tensor Core技术解析（上）
Tensor Core技术解析(上) NVIDIA在SIGGRAPH 2018上正式发布了新一代GPU架构--Turing(图灵),黄仁勋称Turing架构是自2006年CUDA GPU发明以来最大的 ...
用NVIDIA Tensor Cores和TensorFlow 2加速医学图像分割
用NVIDIA Tensor Cores和TensorFlow 2加速医学图像分割 Accelerating Medical Image Segmentation with NVIDIA Tensor ...

Nvidia Tensor Core-MMA PTX编程入门