cuda的ldmatrix指令的详细解释

<font color=red size=6>如果对本篇文章有疑问，欢迎直接评论，我必然秒回！！</font>

• ldmatrix指令是 PTX 级别的指令，它是个warp级别的数据加载指令，当然数据是从shared memory中加载到32个cuda thread中的寄存器中。

1. ldmatrix指令的使用格式例子: ldmatrix.sync.aligned.m8n8.x1.shared.b16 { %0 }, [ %1 ];

• 直接看例子吧，例如这个指令ldmatrix.sync.aligned.m8n8.x1.shared.b16 { %0 }, [ %1 ];
• 这个PTX的指令呢？
  • 他的意思就是让一个warp中的32个线程，从shared memory中加载1(x1)个8*8(m8n8)的矩阵，这个矩阵有8行！
  • 每行的8个元素必须连续存放，不同行之间可以不连续
  • 这个矩阵的元素粒度是b16，也就是一个元素占两个字节。
  • 那也就是说这个warp读取了8*8=64个元素，warp中的每个cuda thread占据了两个元素，因此正好是一个32位的寄存器！
• 由于每行的位置可以不连续，因此用户需要指定8行的开始地址，这个需要用户确保thread0-thread7的%1寄存器填充的是这8个地址即可
  • 至于thread 8-thread 31 的 %1 寄存器数据，那就可以随便设置了！
• 返回的32位数据，也就是在0%寄存器里了。
• 我们好奇的第一个问题就是，这64个元素是如何分布在32个cuda thread之间的呢？
  • 答案就是看下面这个图即可，红色的表示线程id，每个线程被占据了两个数据，也就是正好是1个32位寄存器！
  • 图上只给出了thread0-thread7的情况，其他线程的占据的数据以此类推即可，我这里省略了。
  • 图中一共是8行，注意每行的首地址必须由thread0-thread7的%1寄存器指定!

image.png

2. ldmatrix指令的使用格式例子：ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %0, %1, %2, %3 }, [ %4 ]

• 这个例子使用更广泛，他的意思就是让一个warp中的32个线程，从shared memory中加载4个8*8的矩阵，矩阵的元素仍然是16位！
  • 因此这个时候需要指定32个地址了，也就是thread0-thread31的%4寄存器需要指定这32个地址！
  • 同时每个cuda thread应该瓜分到4*8*8/32=8个bf16元素，也就是4个32位寄存器啦！
  • 所以他的返回值有四个寄存器了！

ldmatrix.sync.aligned指令的例子

• 下面的代码非常简单，值得一看。编译命令是 nvcc A.cu -arch sm_80
• 下面的代码就是让一个warp中的32个线程，从shared memory中加载4个8*8的矩阵，但是我这个里面写的是uint32_t元素类型，因此其实是4个8*4的矩阵！
• 下面代码打印出线程1占据的四个32位寄存器的值！

#include <stdio.h>
#include <iostream>

__global__ void helloFromGPU (void)
{
  __shared__ uint32_t aTile[4*8*4];

  int tidx = threadIdx.x + blockDim.x * threadIdx.y;
  // 下面的代码是把smem中的4*8*4的矩阵，初始化数值！
  if (tidx == 0) {
    for (int i = 0; i < 4*8*4; ++i) {
        aTile[i] = I;
    }
  }
  __syncthreads();

  int aTile_index = tidx % 16 * 8 + tidx / 16 * 4;
  uint32_t a[4];
  uint32_t smem = __cvta_generic_to_shared(aTile+aTile_index);
  asm("ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %0, %1, %2, %3 }, [ %4 ];\n"
  : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) 
  : "r"(smem)
  );

  if (tidx == 1) {
    printf("%d \n", (a[0])); printf("%d \n", (a[1]));
    printf("%d \n", (a[2])); printf("%d \n", (a[3]));
  }
}

int main(void) {
uint3 block = {32,1,1};
uint3 grid = {1,1,1};
helloFromGPU <<<grid, block>>>();

cudaDeviceReset();
return 0;
}

• 上面代码中int aTile_index = tidx % 16 * 8 + tidx / 16 * 4;这个代码是让
• 每个线程分别指向每行的首地址！
• 其中tidx % 16 表示的是行id，
• tidx / 16 * 4 表示的是列id，
  • 从按照b16的角度来看，其实列id应该是tidx / 16 * 8，
  • 但是我代码里面写的是uint32类型，所以就是tidx / 16 * 4了！
• 最后打印出来的四个数字如下图所示哦！

image.png

为什么需要这样的指令呢？

• 为啥需要ldmatrix指令呢？是因为这个指令主要和mma指令搭配使用的！
• 也就是先用ldmatrix指令将数据从shared memory中加载到寄存器，然后调用 mma 指令计算！
• 请我们来看一下这个链接 9.7.13.4.8. Matrix Fragments for mma.m16n8k16 with floating point type，这个链接上展示了mma.m16n8k16指令。
• 这个指令的功能是计算A矩阵16*16和B矩阵16*8，然后得到一个16*8的矩阵C。
• 其中A矩阵16*16的元素分布在32线程的寄存器中，每个线程占据着8个元素，那么这么多元素是如何分布在32个线程中的呢？
• 就是下面这个图片所示哦

image.png

• 从上图可以看出，mma指令对于输入在32个cuda thread之间的分布恰好就是 ldmatrix指令那样！
• B矩阵16*8的元素分布在32个寄存器中，那么这么多元素是如何分布在32个寄存器中的呢？
• 就是下面这个图片所示哦

image.png

• 从上图可以看出，mma指令对于输入在32个cuda thread之间的分布恰好就是 ldmatrix指令那样！
  • 当然啦，这个要求smem中B矩阵必须是col major的哦！，否则就不能调用ldmatrix指令哦！