Nvidia Tensor Core-MMA PTX编程入门

目录

1 PTX (Parallel Thread Execution)

2 MMA (Matrix Multiply Accumulate) PTX

3 LDMATRIX PTX

4 示例

5 底层代码

6 其他

6.1 HGEMM优化


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 % 4

row = groupID      for ai where 0 <= i < 2 || 4 <= i < 6
      groupID + 8  Otherwise

col = (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 % 4

row = (threadID_in_group * 2) + (i & 0x1)      for bi where i < 2 
      (threadID_in_group * 2) + (i & 0x1) + 8  for bi where i >= 2

col = groupID

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

groupID           = %laneid >> 2
threadID_in_group = %laneid % 4

row = groupID      for ci where i < 2
      groupID + 8  for ci where i >= 2

col = (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。这个例子有难度,但不多。源码在cuda_hgemm

#define MMA_M 16
#define MMA_N 8
#define MMA_K 16

#define WARP_SIZE 32

__global__ void mmaNaiveKernel(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 warp_row = blockIdx.y * MMA_M;
    const size_t warp_col = blockIdx.x * MMA_N;

    if (warp_row >= M || warp_col >= N) {
        return;
    }

    __shared__ half A_shmem[MMA_M][MMA_K];
    __shared__ half B_shmem[MMA_N][MMA_K];
    __shared__ half C_shmem[MMA_M][MMA_N];

    const size_t lane_id = threadIdx.x % WARP_SIZE;

    uint32_t RC[2] = {0, 0};

#pragma unroll
    for (size_t i = 0; i < K_tiles; ++i) {
        *((int4 *)(&A_shmem[lane_id / 2][0]) + lane_id % 2) =
            *((int4 *)(&A[(warp_row + lane_id / 2) * K + i * MMA_K]) + lane_id % 2);

        if (lane_id < MMA_N * 2) {
            *((int4 *)(&B_shmem[lane_id / 2][0]) + lane_id % 2) =
                *((int4 *)(&B[i * MMA_K + (warp_col + lane_id / 2) * K]) + lane_id % 2);
        }

        __syncthreads();

        uint32_t RA[4];
        uint32_t RB[2];

        uint32_t A_shmem_lane_addr = __cvta_generic_to_shared(&A_shmem[lane_id % 16][(lane_id / 16) * 8]);
        LDMATRIX_X4(RA[0], RA[1], RA[2], RA[3], A_shmem_lane_addr);

        uint32_t B_shmem_lane_addr = __cvta_generic_to_shared(&B_shmem[lane_id % 8][((lane_id / 8) % 2) * 8]);
        LDMATRIX_X2(RB[0], RB[1], B_shmem_lane_addr);

        HMMA16816(RC[0], RC[1], RA[0], RA[1], RA[2], RA[3], RB[0], RB[1], RC[0], RC[1]);

        __syncthreads();
    }

    *((uint32_t *)(&C_shmem[lane_id / 4][0]) + lane_id % 4) = RC[0];
    *((uint32_t *)(&C_shmem[lane_id / 4 + 8][0]) + lane_id % 4) = RC[1];

    __syncthreads();

    if (lane_id < MMA_M) {
        *((int4 *)(&C[(warp_row + lane_id) * N + warp_col])) = *((int4 *)(&C_shmem[lane_id][0]));
    }
}

void mmaNaive(half *A, half *B, half *C, size_t M, size_t N, size_t K) {
    dim3 block(WARP_SIZE);
    dim3 grid(div_ceil(N, MMA_N), div_ceil(M, MMA_M));

    mmaNaiveKernel<<<grid, block>>>(A, B, C, M, N, K);
}

5 底层代码

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

	Function : _Z14mmaNaiveKernelPK6__halfS1_PS_mmm
	.headerflags    @"EF_CUDA_SM86 EF_CUDA_PTX_SM(EF_CUDA_SM86)"
        /*0000*/                   IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] ;                  /* 0x00000a00ff017624 */
                                                                                            /* 0x000fc400078e00ff */
        /*0010*/                   S2R R2, SR_CTAID.X ;                                     /* 0x0000000000027919 */
                                                                                            /* 0x000e280000002500 */
        /*0020*/                   S2R R0, SR_CTAID.Y ;                                     /* 0x0000000000007919 */
                                                                                            /* 0x000e620000002600 */
        /*0030*/                   IMAD.SHL.U32 R2, R2, 0x8, RZ ;                           /* 0x0000000802027824 */
                                                                                            /* 0x001fe400078e00ff */
        /*0040*/                   IMAD.SHL.U32 R0, R0, 0x10, RZ ;                          /* 0x0000001000007824 */
                                                                                            /* 0x002fc600078e00ff */
        /*0050*/                   ISETP.GE.U32.AND P0, PT, R2, c[0x0][0x180], PT ;         /* 0x0000600002007a0c */
                                                                                            /* 0x000fe40003f06070 */
        /*0060*/                   ISETP.GE.U32.AND P1, PT, R0, c[0x0][0x178], PT ;         /* 0x00005e0000007a0c */
                                                                                            /* 0x000fe40003f26070 */
        /*0070*/                   ISETP.GE.U32.AND.EX P0, PT, RZ, c[0x0][0x184], PT, P0 ;  /* 0x00006100ff007a0c */
                                                                                            /* 0x000fc80003f06100 */
        /*0080*/                   ISETP.GE.U32.OR.EX P0, PT, RZ, c[0x0][0x17c], P0, P1 ;   /* 0x00005f00ff007a0c */
                                                                                            /* 0x000fda0000706510 */
        /*0090*/               @P0 EXIT ;                                                   /* 0x000000000000094d */
                                                                                            /* 0x000fea0003800000 */
        /*00a0*/                   S2R R9, SR_TID.X ;                                       /* 0x0000000000097919 */
                                                                                            /* 0x000e220000002100 */
        /*00b0*/                   IMAD.MOV.U32 R10, RZ, RZ, c[0x0][0x188] ;                /* 0x00006200ff0a7624 */
                                                                                            /* 0x000fe200078e00ff */
        /*00c0*/                   ULDC.64 UR4, c[0x0][0x118] ;                             /* 0x0000460000047ab9 */
                                                                                            /* 0x000fe20000000a00 */
        /*00d0*/                   IMAD.MOV.U32 R5, RZ, RZ, 0x4 ;                           /* 0x00000004ff057424 */
                                                                                            /* 0x000fe200078e00ff */
        /*00e0*/                   CS2R R16, SRZ ;                                          /* 0x0000000000107805 */
                                                                                            /* 0x000fe2000001ff00 */
        /*00f0*/                   IMAD.MOV.U32 R8, RZ, RZ, c[0x0][0x18c] ;                 /* 0x00006300ff087624 */
                                                                                            /* 0x000fe200078e00ff */
        /*0100*/                   LOP3.LUT P0, RZ, R10.reuse, 0xf, RZ, 0xc0, !PT ;         /* 0x0000000f0aff7812 */
                                                                                            /* 0x040fe4000780c0ff */
        /*0110*/                   SHF.R.U64 R10, R10, R5, c[0x0][0x18c] ;                  /* 0x000063000a0a7619 */
                                                                                            /* 0x000fe40000001205 */
        /*0120*/                   LOP3.LUT P0, RZ, RZ, c[0x0][0x18c], RZ, 0xc0, P0 ;       /* 0x00006300ffff7a12 */
                                                                                            /* 0x000fc4000000c0ff */
        /*0130*/                   SHF.R.U32.HI R8, RZ, 0x4, R8 ;                           /* 0x00000004ff087819 */
                                                                                            /* 0x000fe40000011608 */
        /*0140*/                   SEL R7, RZ, 0xffffffff, !P0 ;                            /* 0xffffffffff077807 */
                                                                                            /* 0x000fc80004000000 */
        /*0150*/                   ISETP.NE.U32.AND P1, PT, R10, R7, PT ;                   /* 0x000000070a00720c */
                                                                                            /* 0x000fc80003f25070 */
        /*0160*/                   ISETP.NE.AND.EX P1, PT, R8, R7, PT, P1 ;                 /* 0x000000070800720c */
                                                                                            /* 0x000fe40003f25310 */
        /*0170*/                   LOP3.LUT R19, R9.reuse, 0x1f, RZ, 0xc0, !PT ;            /* 0x0000001f09137812 */
                                                                                            /* 0x041fe200078ec0ff */
        /*0180*/                   IMAD.SHL.U32 R5, R9, 0x4, RZ ;                           /* 0x0000000409057824 */
                                                                                            /* 0x000fc600078e00ff */
        /*0190*/                   ISETP.GT.U32.AND P2, PT, R19.reuse, 0xf, PT ;            /* 0x0000000f1300780c */
                                                                                            /* 0x040fe20003f44070 */
        /*01a0*/                   IMAD.SHL.U32 R4, R19, 0x4, RZ ;                          /* 0x0000000413047824 */
                                                                                            /* 0x000fe200078e00ff */
        /*01b0*/                   LOP3.LUT R5, R5, 0xc, RZ, 0xc0, !PT ;                    /* 0x0000000c05057812 */
                                                                                            /* 0x000fc800078ec0ff */
        /*01c0*/                   LOP3.LUT R4, R4, 0x70, RZ, 0xc0, !PT ;                   /* 0x0000007004047812 */
                                                                                            /* 0x000fca00078ec0ff */
        /*01d0*/                   IMAD.IADD R18, R4, 0x1, R5 ;                             /* 0x0000000104127824 */
                                                                                            /* 0x000fe200078e0205 */
        /*01e0*/              @!P1 BRA 0x580 ;                                              /* 0x0000039000009947 */
                                                                                            /* 0x000fea0003800000 */
        /*01f0*/                   SHF.R.U64 R13, R19, 0x1, RZ ;                            /* 0x00000001130d7819 */
                                                                                            /* 0x000fe200000012ff */
        /*0200*/                   IMAD.SHL.U32 R4, R9, 0x10, RZ ;                          /* 0x0000001009047824 */
                                                                                            /* 0x000fe200078e00ff */
        /*0210*/                   SEL R20, RZ, 0x1, !P0 ;                                  /* 0x00000001ff147807 */
                                                                                            /* 0x000fe20004000000 */
        /*0220*/                   IMAD.SHL.U32 R5, R19, 0x8, RZ ;                          /* 0x0000000813057824 */
                                                                                            /* 0x000fe200078e00ff */
        /*0230*/                   IADD3 R7, P3, R2, R13.reuse, RZ ;                        /* 0x0000000d02077210 */
                                                                                            /* 0x080fe20007f7e0ff */
        /*0240*/                   CS2R R16, SRZ ;                                          /* 0x0000000000107805 */
                                                                                            /* 0x000fe2000001ff00 */
        /*0250*/                   IADD3 R11, P1, R0, R13, RZ ;                             /* 0x0000000d000b7210 */
                                                                                            /* 0x000fe40007f3e0ff */
        /*0260*/                   LOP3.LUT R21, R4, 0x10, RZ, 0xc0, !PT ;                  /* 0x0000001004157812 */
                                                                                            /* 0x000fe200078ec0ff */
        /*0270*/                   IMAD.X R6, RZ, RZ, RZ, P3 ;                              /* 0x000000ffff067224 */
                                                                                            /* 0x000fe200018e06ff */
        /*0280*/                   LOP3.LUT R4, R5, 0x8, RZ, 0xc0, !PT ;                    /* 0x0000000805047812 */
                                                                                            /* 0x000fe200078ec0ff */
        /*0290*/                   IMAD.X R12, RZ, RZ, RZ, P1 ;                             /* 0x000000ffff0c7224 */
                                                                                            /* 0x000fe200008e06ff */
        /*02a0*/                   IADD3 R20, P0, R20, R10, RZ ;                            /* 0x0000000a14147210 */
                                                                                            /* 0x000fe20007f1e0ff */
        /*02b0*/                   IMAD R6, R6, c[0x0][0x188], RZ ;                         /* 0x0000620006067a24 */
                                                                                            /* 0x000fc400078e02ff */
        /*02c0*/                   IMAD.MOV.U32 R5, RZ, RZ, RZ ;                            /* 0x000000ffff057224 */
                                                                                            /* 0x000fe400078e00ff */
        /*02d0*/                   IMAD R12, R12, c[0x0][0x188], RZ ;                       /* 0x000062000c0c7a24 */
                                                                                            /* 0x000fe400078e02ff */
        /*02e0*/                   IMAD.SHL.U32 R10, R9, 0x20, RZ ;                         /* 0x00000020090a7824 */
                                                                                            /* 0x000fe400078e00ff */
        /*02f0*/                   IMAD R27, R7.reuse, c[0x0][0x18c], R6 ;                  /* 0x00006300071b7a24 */
                                                                                            /* 0x040fe400078e0206 */
        /*0300*/                   IMAD.WIDE.U32 R6, R7, c[0x0][0x188], R4.reuse ;          /* 0x0000620007067a25 */
                                                                                            /* 0x100fe200078e0004 */
        /*0310*/                   LOP3.LUT R23, R10.reuse, 0x1e0, RZ, 0xc0, !PT ;          /* 0x000001e00a177812 */
                                                                                            /* 0x040fe400078ec0ff */
        /*0320*/                   LOP3.LUT R25, R10, 0xe0, RZ, 0xc0, !PT ;                 /* 0x000000e00a197812 */
                                                                                            /* 0x000fe200078ec0ff */
        /*0330*/                   IMAD.SHL.U32 R9, R9, 0x2, RZ ;                           /* 0x0000000209097824 */
                                                                                            /* 0x000fe200078e00ff */
        /*0340*/                   LEA R24, P1, R6, c[0x0][0x168], 0x1 ;                    /* 0x00005a0006187a11 */
                                                                                            /* 0x000fe200078208ff */
        /*0350*/                   IMAD.WIDE.U32 R4, R11, c[0x0][0x188], R4 ;               /* 0x000062000b047a25 */
                                                                                            /* 0x000fc600078e0004 */
        /*0360*/                   LOP3.LUT R10, R9, 0x10, RZ, 0xc0, !PT ;                  /* 0x00000010090a7812 */
                                                                                            /* 0x000fe200078ec0ff */
        /*0370*/                   IMAD R31, R11, c[0x0][0x18c], R12 ;                      /* 0x000063000b1f7a24 */
                                                                                            /* 0x000fe200078e020c */
        /*0380*/                   LOP3.LUT R12, R19, 0x10, RZ, 0xc0, !PT ;                 /* 0x00000010130c7812 */
                                                                                            /* 0x000fe200078ec0ff */
        /*0390*/                   IMAD.IADD R27, R7, 0x1, R27 ;                            /* 0x00000001071b7824 */
                                                                                            /* 0x000fe200078e021b */
        /*03a0*/                   LEA R22, P3, R4, c[0x0][0x160], 0x1 ;                    /* 0x0000580004167a11 */
                                                                                            /* 0x000fe200078608ff */
        /*03b0*/                   IMAD.IADD R31, R5, 0x1, R31 ;                            /* 0x00000001051f7824 */
                                                                                            /* 0x000fe400078e021f */
        /*03c0*/                   IMAD R21, R13, 0x20, R21 ;                               /* 0x000000200d157824 */
                                                                                            /* 0x000fe200078e0215 */
        /*03d0*/                   LEA.HI.X R27, R6, c[0x0][0x16c], R27, 0x1, P1 ;          /* 0x00005b00061b7a11 */
                                                                                            /* 0x000fe200008f0c1b */
        /*03e0*/                   IMAD.IADD R23, R23, 0x1, R12 ;                           /* 0x0000000117177824 */
                                                                                            /* 0x000fe200078e020c */
        /*03f0*/                   LEA.HI.X R31, R4, c[0x0][0x164], R31, 0x1, P3 ;          /* 0x00005900041f7a11 */
                                                                                            /* 0x000fe200018f0c1f */
        /*0400*/                   IMAD.IADD R25, R25, 0x1, R10 ;                           /* 0x0000000119197824 */
                                                                                            /* 0x000fc400078e020a */
        /*0410*/                   IMAD.X R30, RZ, RZ, R8, P0 ;                             /* 0x000000ffff1e7224 */
                                                                                            /* 0x000fc400000e0608 */
        /*0420*/                   ISETP.GT.U32.AND P0, PT, R19, 0xf, PT ;                  /* 0x0000000f1300780c */
                                                                                            /* 0x000fe20003f04070 */
        /*0430*/                   IMAD.MOV.U32 R4, RZ, RZ, R22 ;                           /* 0x000000ffff047224 */
                                                                                            /* 0x000fe400078e0016 */
        /*0440*/                   IMAD.MOV.U32 R5, RZ, RZ, R31 ;                           /* 0x000000ffff057224 */
                                                                                            /* 0x000fcc00078e001f */
        /*0450*/                   LDG.E.128.CONSTANT R4, [R4.64] ;                         /* 0x0000000404047981 */
                                                                                            /* 0x000ea8000c1e9d00 */
        /*0460*/              @!P0 IMAD.MOV.U32 R26, RZ, RZ, R24 ;                          /* 0x000000ffff1a8224 */
                                                                                            /* 0x000fca00078e0018 */
        /*0470*/              @!P0 LDG.E.128.CONSTANT R8, [R26.64] ;                        /* 0x000000041a088981 */
                                                                                            /* 0x0000e2000c1e9d00 */
        /*0480*/                   IADD3 R22, P1, R22, 0x20, RZ ;                           /* 0x0000002016167810 */
                                                                                            /* 0x000fe40007f3e0ff */
        /*0490*/                   IADD3 R24, P3, R24, 0x20, RZ ;                           /* 0x0000002018187810 */
                                                                                            /* 0x000fc60007f7e0ff */
        /*04a0*/                   IMAD.X R31, RZ, RZ, R31, P1 ;                            /* 0x000000ffff1f7224 */
                                                                                            /* 0x000fe400008e061f */
        /*04b0*/                   IMAD.X R27, RZ, RZ, R27, P3 ;                            /* 0x000000ffff1b7224 */
                                                                                            /* 0x001fe200018e061b */
        /*04c0*/                   STS.128 [R21], R4 ;                                      /* 0x0000000415007388 */
                                                                                            /* 0x004fe80000000c00 */
        /*04d0*/              @!P0 STS.128 [R21+0x200], R8 ;                                /* 0x0002000815008388 */
                                                                                            /* 0x008fe80000000c00 */
        /*04e0*/                   BAR.SYNC 0x0 ;                                           /* 0x0000000000007b1d */
                                                                                            /* 0x000fec0000000000 */
        /*04f0*/                   LDSM.16.M88.2 R28, [R25+0x200] ;                         /* 0x00020000191c783b */
                                                                                            /* 0x020fe80000000100 */
        /*0500*/                   LDSM.16.M88.4 R12, [R23] ;                               /* 0x00000000170c783b */
                                                                                            /* 0x000e280000000200 */
        /*0510*/                   BAR.SYNC 0x0 ;                                           /* 0x0000000000007b1d */
                                                                                            /* 0x000fec0000000000 */
        /*0520*/                   IADD3 R20, P0, R20, -0x1, RZ ;                           /* 0xffffffff14147810 */
                                                                                            /* 0x000fc80007f1e0ff */
        /*0530*/                   IADD3.X R30, R30, -0x1, RZ, P0, !PT ;                    /* 0xffffffff1e1e7810 */
                                                                                            /* 0x000fe400007fe4ff */
        /*0540*/                   ISETP.NE.U32.AND P0, PT, R20, RZ, PT ;                   /* 0x000000ff1400720c */
                                                                                            /* 0x000fc80003f05070 */
        /*0550*/                   ISETP.NE.AND.EX P0, PT, R30, RZ, PT, P0 ;                /* 0x000000ff1e00720c */
                                                                                            /* 0x000fe20003f05300 */
        /*0560*/                   HMMA.16816.F16 R16, R12, R28, R16 ;                      /* 0x0000001c0c10723c */
                                                                                            /* 0x001b580000000810 */
        /*0570*/               @P0 BRA 0x420 ;                                              /* 0xfffffea000000947 */
                                                                                            /* 0x000fca000383ffff */
        /*0580*/                   NOP ;                                                    /* 0x0000000000007918 */
                                                                                            /* 0x000fcc0000000000 */
        /*0590*/                   STS [R18+0x300], R16 ;                                   /* 0x0003001012007388 */
                                                                                            /* 0x020fe80000000800 */
        /*05a0*/                   STS [R18+0x380], R17 ;                                   /* 0x0003801112007388 */
                                                                                            /* 0x000fe80000000800 */
        /*05b0*/                   BAR.SYNC 0x0 ;                                           /* 0x0000000000007b1d */
                                                                                            /* 0x000fec0000000000 */
        /*05c0*/               @P2 EXIT ;                                                   /* 0x000000000000294d */
                                                                                            /* 0x000fea0003800000 */
        /*05d0*/                   LEA R4, R19, 0x300, 0x4 ;                                /* 0x0000030013047811 */
                                                                                            /* 0x000fe200078e20ff */
        /*05e0*/                   IMAD.MOV.U32 R3, RZ, RZ, RZ ;                            /* 0x000000ffff037224 */
                                                                                            /* 0x000fe200078e00ff */
        /*05f0*/                   IADD3 R19, P0, R0, R19, RZ ;                             /* 0x0000001300137210 */
                                                                                            /* 0x000fc80007f1e0ff */
        /*0600*/                   LDS.128 R4, [R4] ;                                       /* 0x0000000004047984 */
                                                                                            /* 0x000e220000000c00 */
        /*0610*/                   IMAD.X R0, RZ, RZ, RZ, P0 ;                              /* 0x000000ffff007224 */
                                                                                            /* 0x000fe400000e06ff */
        /*0620*/                   IMAD.WIDE.U32 R2, R19, c[0x0][0x180], R2 ;               /* 0x0000600013027a25 */
                                                                                            /* 0x000fc800078e0002 */
        /*0630*/                   IMAD R0, R0, c[0x0][0x180], RZ ;                         /* 0x0000600000007a24 */
                                                                                            /* 0x000fe200078e02ff */
        /*0640*/                   LEA R8, P0, R2, c[0x0][0x170], 0x1 ;                     /* 0x00005c0002087a11 */
                                                                                            /* 0x000fc600078008ff */
        /*0650*/                   IMAD R19, R19, c[0x0][0x184], R0 ;                       /* 0x0000610013137a24 */
                                                                                            /* 0x000fc800078e0200 */
        /*0660*/                   IMAD.IADD R3, R3, 0x1, R19 ;                             /* 0x0000000103037824 */
                                                                                            /* 0x000fca00078e0213 */
        /*0670*/                   LEA.HI.X R9, R2, c[0x0][0x174], R3, 0x1, P0 ;            /* 0x00005d0002097a11 */
                                                                                            /* 0x000fca00000f0c03 */
        /*0680*/                   STG.E.128 [R8.64], R4 ;                                  /* 0x0000000408007986 */
                                                                                            /* 0x001fe2000c101d04 */
        /*0690*/                   EXIT ;                                                   /* 0x000000000000794d */
                                                                                            /* 0x000fea0003800000 */
        /*06a0*/                   BRA 0x6a0;                                               /* 0xfffffff000007947 */
                                                                                            /* 0x000fc0000383ffff */
        /*06b0*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*06c0*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*06d0*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*06e0*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*06f0*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0700*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0710*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0720*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0730*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0740*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0750*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0760*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
        /*0770*/                   NOP;                                                     /* 0x0000000000007918 */
                                                                                            /* 0x000fc00000000000 */
		..........

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

6 其他

6.1 HGEMM优化

与WMMA API类似,学习MMA PTX的目标在于调用Tensor Core优化HGEMM,相比于cublas,MMA的性能究竟如何?我们后续再聊。可以参考开源代码cuda_hgemm

  • 4
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值