[CUDA] ptx使用笔记

官方网站

PTX ISA doc: 包括PTX的一些指令集的使用手册。
PTX compiler API: The PTX Compiler APIs are a set of APIs which can be used to compile a PTX program into GPU assembly code. 一个可能被使用的场景是在分离的ptx与主程序应用中，通过更新ptx模块代码，让主程序使用最新的ptx特性。（而不是将ptx嵌入到主程序中，然后每次都有编译所有的程序。）
PTX Writer’s Guide to Interoperability互操作性: This document defines the Application Binary Interface (ABI) for the CUDA® architecture when generating PTX. By following the ABI, external developers can generate compliant PTX code that can be linked with other code.
Inline PTX Assembly in CUDA: The reference guide for inlining PTX (parallel thread execution) assembly statements into CUDA. 也就是怎么在cuda中写PTX代码。

常用指令

1. 基本写法

PTX:

• 源模块是ASCII文本。行由换行字符（\n）分隔。
• PTX是区分大小写的，并使用小写字母作为关键字。
• PTX源模块具有汇编语言风格的语法，包括指令操作码和操作数。伪操作用于指定符号和地址管理。
• 每个PTX模块必须以.version指令开始，指定PTX语言版本，然后是.target指令，指定假定的目标架构。
• 指令关键字以点开头，因此不可能会与用户定义的标识符发生冲突

// 声明一个寄存器变量addr， 类型是u64.
.reg .u64 addr; 

2. 一些基本指令用法

• 包括: .pred, setp, @p, ld.global.v4.u32 ... 以及cuda中如何嵌入asm代码

参考实例来源：https://github.com/mit-han-lab/torchsparse/blob/master/torchsparse/backend/utils/memory.cuh

template <int bytes>
struct global_load;

template <>
struct global_load<16> {
   
  __device__ __inline__ global_load(uint4& D, void const* ptr, int pred_guard) {
   
    uint4& data = *reinterpret_cast<uint4*>(&D);
    // 应该也可以使用 __asm__ __volatile__.
    // 下面这段ptx指令表示为，但是更加有效：
    // if (static_cast<int>(pred_guard & (1 << ldg_idx)) != 0) {
   
    //  data = *(ptr_ldg + ldg_idx);
    // },。
    asm volatile(
        "{\n"
        // 下面一段话代表声明一个谓词变量， 谓词变量的应该可以等同于表达式或者一个行为变量。
        "  .reg .pred p;\n"
        // .pred的变量经常是和setp一起使用的，这ptx的意思是将p = (int)(pred_guard & 1) != 0
        // setp:  Comparison and Selection Instructions： Compare two numeric values with a relational operator, and (optionally) combine this result with a predicate value by applying a Boolean operator.
        // ne 表示不等于not equal.
        // %5 表示ascii后面的第五个参数，从0开始计数； 所以是(int)(pred_guard & 1)的值
        "  setp.ne.b32 p, %5, 0;\n"
        // 注意下面这四行命令，表示将data.x = data.x, 因为D声明为 uint4 D = make_uint4(0,0,0,0);
        // 所以这四行命令其实是延迟执行这个初始化，也就是将D.x, D.y,D.z,D.w 初始化为0;
        "  mov.b32 %0, %6;\n"
        "  mov.b32 %1, %7;\n"
        "  mov.b32 %2, %8;\n"
        "  mov.b32 %3, %9;\n"
        // @p 表示if(p==true) / if(p)的意思； 如果p=true, 则执行ld.global.v4
        // 注意v4表示vector为4 elems。 写法固定，用大括号括起来目的地址，用中括号括起来源地址。
        "  @p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n"
        "}\n"
        // =r 中有等号，表示目的地址； 而后面没有等号且用“：”分开的是原地址； 注意目的地址和原地址要用":"分开。
        : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w)
        : "l"(ptr), "r"((int)(pred_guard & 1)), "r"(data.x), "r"(data.y),
          "r"(data.z), "r"(data.w));
  }
};

3. mma指令与ld.matrix

注意使用cuda的tensor core编写矩阵乘的时候，有两种方式，一种是上层函数wmma来实现矩阵分块使用tensocore加速； 一种是底层ptx的mma指令来实现；两者区别是：

• wmma 操作基本上是将一个warp的数据全部导入，不够灵活，对支持密集数据友好。
  而mma操作是支持稀疏数据的，需要显性定义每个thread导入数据的过程。

template <int K_tile>
global void __launch_bounds__(128)
    ConvForSingleKernel(int M, int K_original, int N, int j_factor,
                        half* restrict input, half* restrict kernel,
                        half* restrict output);

// 注意每个warp要使用不同的shared 区段，如果使用同一区段会存在一些问题，因为不同
// warp的执行先后不一样，如果存在threadIdx.x==0这样的判断，导致threadIdx.y >
// 0的时候 还没执行完，
// 但是threadIdx.y==0的时候已经执行了新的数据copy到B_shared，这就会导致问题
template <>
global void __launch_bounds__(128) ConvForSingleKernel<8>(
    int M, int K_original, int N, int j_factor, half* restrict input,
    half* restrict kernel, half* restrict output) {
   
  const int K_tile = 8;
  shared half A_shared[2048];
  shared half B_shared[1024];
  float C_warp[K_tile];
  // shape(32, 8) is stored into 32 * 8 registers.
  half A_shared_warp[K_tile];
  // shape(8, 8) is stored into 32 * 2 registers.
  half B_shared_warp[2];
  for (int i = 0; i < K_tile; ++i) {
   
    C_warp[i] = 0.0;
  }
  half* A_ptr = input + (blockIdx.x / j_factor) * K_original * 128 +
                threadIdx.y * WARP_SIZE * K_original;
  half* B_ptr = kernel + (blockIdx.x % j_factor) * K_tile;
  half* C_ptr = output + (blockIdx.x % j_factor) * K_tile;
  for (int a_ax_1 = 0; a_ax_1 < (K_original / K_tile); ++a_ax_1) {
   
    if (threadIdx.x == 0) {
   
#pragma unroll
      for (int b_ax_0 = 0; b_ax_0 < K_tile; ++b_ax_0) {
   
        // +threadIdx.y * K_tile * K_tile
        *reinterpret_cast<uint4*>(B_shared + b_ax_0 * K_tile +
                                  a_ax_1 * K_tile * K_tile) =
            *reinterpret_cast<uint4*>(B_ptr + b_ax_0 * N + a_ax_1 * K_tile * N);
      }
    }
  }
  __syncthreads();

  for (int a_ax_1 = 0; a_ax_1 < (K_original / K_tile); ++a_ax_1) {
   
    // load B
    //     if (threadIdx.x == 0) {
   
    // #pragma unroll
    //       for (int b_ax_0 = 0; b_ax_0 < K_tile; ++b_ax_0) {
   
    //         // +threadIdx.y * K_tile * K_tile
    //         *reinterpret_cast<uint4*>(B_shared + b_ax_0 * K_