a hgemm tvm schedule

Aug 31 2022 Technical an hour read (About 6881 words)

这里记录的是我想从tvm的tensor expression出发，参考一下cutlass efficient gemm的思路，一步一步优化一下GEMM的一些思考和碎碎念，目的是为了理解cutlass优化gemm的思路。

我们使用CUTLASS Profiler来运行一个gemm的运算，并用nsight compute dump下来其运行过程中的一些情况，可以拿到他的一些信息，如grid的大小与block的大小等。比如对于16384的float32类型数据的gemm，cutlass的grid size是（512, 16, 1）-> 8192个block， block size是（256，1，1），一共是2,097,152个线程，因为最后产生C的大小是（16384，16384），所以平均每个thread需要产生128个C的元素，结合这些参数的信息，使用tvm的te进行schedule（其实可以试试tensor ir），最后成功打到了和cublas，cutlass相近的性能。

测试GPU: rtx 3090 24GB

CUDA Version: 11.1

TVM Version: 10.0

代码放在：https://github.com/LeiWang1999/tvm_gpu_gemm

naive gemm

最navie的情况，一个thread处理一个像素点的数据，这里我们选大小为1024的gemm为例子(因为这么做太慢了，大shape要跑很久)

for m in range 1024:
    for n in range 1024
        c[m][n] = 0.0
        for k in range 1024:
            c[m][n] += a[m][k]*a[k][n]

对应到tvm的表达式：

    # graph
    nn = 1024
    n = te.var("n")
    n = tvm.runtime.convert(nn)
    m, l = n, n
    A = te.placeholder((l, n), dtype=_dtype, name="A")
    B = te.placeholder((l, m), dtype=_dtype,name="B")
    k = te.reduce_axis((0, l), name="k")
    C = te.compute((m, n), lambda ii, jj: te.sum(A[k, jj] * B[k, ii], axis=k), name="C")

tvm可以显式的bind thread，例如我们把第一个维度的loop bind到blockIdx.x，第二个维度的loop bind到threadIdx.x，就完成了每个thread计算一个C的像素。

    block_x = te.thread_axis("blockIdx.x")
    block_y = te.thread_axis("blockIdx.y")
    thread_x = te.thread_axis("threadIdx.x")
    thread_y = te.thread_axis("threadIdx.y")

    i, j = s[C].op.axis

    s[C].bind(i, block_x)
    s[C].bind(j, thread_x)

运行下来的速度是2.xms是cutlass的二十倍左右，生成的cuda kernel如下：

extern "C" __global__ void __launch_bounds__(1024) default_function_kernel0(float* __restrict__ C, float* __restrict__ A, float* __restrict__ B) {
  C[((((int)blockIdx.x) * 1024) + ((int)threadIdx.x))] = 0.000000e+00f;
  for (int k = 0; k < 1024; ++k) {
    C[((((int)blockIdx.x) * 1024) + ((int)threadIdx.x))] = (C[((((int)blockIdx.x) * 1024) + ((int)threadIdx.x))] + (A[((k * 1024) + ((int)threadIdx.x))] * B[((k * 1024) + ((int)blockIdx.x))]));
  }
}

Blocked Gemm

分析上述生成的代码，每个thread需要从global memory读取两个数据，A和B，而显然，A和B中的数据都是可以被复用的，这就导致同一个数据被使用了几次，就从global memory中读取了几次，所需要地总的读取Global Memory的次数为 M*N*K*2 次，写M*N次。

上图是CUTLASS实现高效GEMM的示意图，cutlass在第一层使用矩阵分块的方式来减少global memory的访问次数，在CPU上优化gemm同样也需要用到矩阵分块，不过cpu上的矩阵分块主要是为了fit L2 Cache的大小，这里的却略有不同，cutlass的图描述地不是很明了，绿色的方块代表的是一个block，负责产出一个C子矩阵的数据，可以通过A的浅蓝色的矩阵部分与B的浅黄色的矩阵部分乘累加得到，因为一个block内的thread共享share memory，所以可以提前把这两部分浅色的矩阵cache到share memory中。

现在暂时不考虑thread block tile这部分内容，则每个block内的thread各自负责一个计算像素，比如，我们一个block负责产生（32，32）个像素则在16384的gemm场景下：

M,K,N: 16384，16384，16384

grid_size: 16384, 16

block_size: 1024 (32,32)

对于grid来说，每一个block需要处理一个小的矩阵，具体到tvm的调度实现代码是：

     bx, xi = s[C].split(C.op.axis[0], factor=(block_h))
    write_code(str(tvm.lower(s, [A, B, C], simple_mode=True)), "progress/2.split_i.cu")
    by, yi = s[C].split(C.op.axis[1], factor=(block_w))
    write_code(str(tvm.lower(s, [A, B, C], simple_mode=True)), "progress/3.split_j.cu")

    s[C].bind(bx, block_x)
    s[C].bind(by, block_y)
    s[C].reorder(bx, by, xi, yi)
    write_code(str(tvm.lower(s, [A, B, C], simple_mode=True)), "progress/4.bind_block.cu")

    s[C].bind(xi, thread_x)
    s[C].bind(yi, thread_y)
    write_code(str(tvm.lower(s, [A, B, C], simple_mode=True)), "progress/5.bind_thread.cu")

split一般会用到factor和nparts这两个参数，其中factor指的是拆分因数，及block_h是拆分之后的内循环，nparts是拆成多少份，及指定的参数是拆分之后的外循环。

把i,j两层循环各自拆分成（_, 32）,(_, 32), 再将这两层循环外提，就可以达到矩阵分块计算的效果，直接从global memory读取数据进行计算，一次计算的时间是4849.64ms，比cutlass要慢十倍。

接下来，给A、B、C加上cache，先以C举例，”local“表示结果暂存到寄存器中，则数据会先写到CC，再写回global：

    CC = s.cache_write(C, "local")

lower:

@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {A: Buffer(A_2: Pointer(float32), float32, [268435456], []),
             B: Buffer(B_2: Pointer(float32), float32, [268435456], []),
             C: Buffer(C_2: Pointer(float32), float32, [268435456], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C}
  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [16384, 16384], []), B_1: B_3: Buffer(B_2, float32, [16384, 16384], []), C_1: C_3: Buffer(C_2, float32, [16384, 16384], [])} {
  allocate(C.local: Pointer(local float32), float32, [1]), storage_scope = local {
    C.local_1: Buffer(C.local, float32, [1], [], scope="local", align=4)[0] = 0f32
    for (k: int32, 0, 16384) {
      let cse_var_1: int32 = (k*16384)
      C.local_1[0] = (C.local_1[0] + (A[((cse_var_1 + (blockIdx.y: int32*32)) + threadIdx.y: int32)]*B[((cse_var_1 + (blockIdx.x: int32*32)) + threadIdx.x: int32)]))
    }
    attr [IterVar(blockIdx.x, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 512;
    attr [IterVar(blockIdx.y, (nullptr), "ThreadIndex", "blockIdx.y")] "thread_extent" = 512;
    attr [IterVar(threadIdx.x, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32;
    attr [IterVar(threadIdx.y, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 32;
    C[((((blockIdx.x*524288) + (threadIdx.x*16384)) + (blockIdx.y*32)) + threadIdx.y)] = C.local_1[0]
  }
}

发现计算都在c_local上进行了，但是thread和block都bind在C上，这样直接生成cuda程序会有错误，我们要把c_local挪到C的第一个维度（因为都已经绑定到thread了，所以这里C的计算没有循环，需要绑定到yi。

s[CC].compute_at(s[C], yi)

这样可以生成cuda程序，验证性能，不过只把C缓存到local，对性能没有影响，估计是NVCC自己优化掉了，缓存A和B才是关键，但是缓存了A和B，就会尴尬地发现，share memory的大小不够用了，对于每个block，需要load两个16384，32大小的子矩阵，以单精度浮点数为例子，一共是256KB，一个SM也就128KB的share memory资源，所以不可以像这样一个thread计算一个像素。

RTX 3090
register file per SM	256 KB
Maximum number of resident threads per SM	1536
Max number of blocks per SM	16
shared memory size per SM	128 KB

Thread Block Tile

这一步需要进一步的决定，每个block里的thread是怎么计算的，对应cutlass efficent gemm中的第二张图，刚才的方法中，一个thread负责计算一个像素点的数据，在计算数据的时候每个thread又需要从global memory中读取两个 16384, 1 的数据，复用性极差。我们只能再进行分块，一次从global memory中load一部分数据进行计算。

在这之后有两种计算方法，第一种是把浅色部分的矩阵按照inner product的方式，及下图所示：

inner product

还有一种是outer product，将最内层的K循环提出到最外层：

对应到cutlass的示意图，很明显，我看了很多加速gemm的资料，大家用的都是第二种外积的方式，不过也没有人说为什么要这么用，第一种内积的方式就是把A的浅色矩阵横过来切，我的理解是，这样算下来有几个不方便的地方，一是外积的方式只需要决定Bk的大小，而外积的方式如果把K固定，那Bm和Bn就有两个参数需要调整，比较麻烦；二是在计算中，总有一些数据会被load多次，造成了额外的访存。虽然这种方式不能缓解inner product的局部性问题，但是在做矩阵分块的时候确实能更好地利用数据，不过外积的方式每次迭代存储的中间结果是一整个矩阵，内积是一小部分矩阵。

More sophisticated algorithms construct the product hierarchically from multiplications of smaller matrix blocks [7]. By computing outer products on small blocks of the input and output matrices, we can more effectively exploit spatial locality and data reuse. Figure 4 illustrates a blocking scheme with parameters I0, J0, and K0.

cite from : https://patterns.eecs.berkeley.edu/?page_id=158

通过这种方式，可以大大缓解之前从global memory获取数据带来的开销，在block计算每个小C矩阵的的时候，需要从global memory读取数据的次数变成了，设小C矩阵的长宽为Bm,Bn, 则总需要的访存量为：
$$
\frac{MN}{b_mb_n}(Kb_m + Kbn) =MNK(\frac{1}{b_m} + \frac{1}{b_n})
$$
也就是说，每次block只需要load BM*BK + BK*BN 大小的矩阵进入share memory，再循环K / BK次以完成整个矩阵乘法。

这样，就有几个超参数需要我们确定，Block size需要设置为多大？一个Block计算多少个数据？下面就这个问题讨论一下。

首先，关于block size的大小，一般会有两个具体的限制作约束，Maximum number of resident blocks per SM 和 Maximum number of resident threads per SM，也就是 SM 上最大同时执行的 block 数量和线程数量。要到达这个目的有多种方法，其中一个最简单的方法是让尽量多的线程同时在 SM 上执行，SM 上并发执行的线程数和SM 上最大支持的线程数的比值，被称为 Occupancy，更高的 Occupancy 代表潜在更高的性能。

显然，一个 kernel 的 block_size 应大于 SM 上最大线程数和最大 block 数量的比值，否则就无法达到 100% 的 Occupancy，对应不同的架构，这个比值不相同.>

对于 V100 、 A100、 GTX 1080 Ti 是 2048 / 32 = 64，对于 RTX 3090 是 1536 / 16 = 96，所以为了适配主流架构，如果静态设置 block_size 不应小于 96。虑到 block 调度的原子性，那么 block_size 应为 SM 最大线程数的约数，否则也无法达到 100% 的 Occupancy，主流架构的 GPU 的 SM 最大线程数的公约是 512，96 以上的约数还包括 128 和 256，也就是到目前为止，block_size 的可选值仅剩下 128 / 256 / 512 三个值。

cite from 如何设置CUDA Kernel中的grid_size和block_size？

不难发现，cutlass里的绝大多数kernel都选用的是128和256这两个值，当block中的thread数量变多，单个thread能吃到的资源就变小，反而会对编程产生限制。

其次，关于一个Block计算多少个数据，由于我们已经确定了block的大小，即一个block里可以有128/256个thread，乘上每个thread处理的像素个数，就可以得到每个block要算几个像素了，从刚才得到的公式中可以知道，一个block处理的像素点越多的时候，他的访存量就会越小，但显然这个值是有上限的，一个sm上的thread会均摊share memory。

有大概下面几个约束：

每个block至少需要往share memory里塞 (BM * BK + BK * BN)个数据，如果要乒乓缓存还需要double，一个SM上可以同时处理多个Block（貌似还是碎片化管理，动静态分配结合的方式，所以这也没有一个固定的最大值，可能要用心去感受）。
每个block有最大线程数的限制，如上表所示，3090的一个block的最大线程数是1536。

所以，BM、BN、BK、blockIdx的大小都是可以tune的，这些值确定下来，加上所要处理的所有像素就可以算出grid size。

回归到每个block里的thread是如何计算的，图二一个Block里有1024个线程，包含了八个warp tile，一个warp Tile包含了四个warp，一个warp有32个线程。

对与fp32的cutlass，测了一下在cudacore上最好的square 16384 kernel是cutlass_simt_sgemm_128x128_8x2_nt_align1，grid_size是（1024，16，1），block size是（256，1，1），平均一个线程处理六十四个数据，一个block处理16384个数据。

grid_size是（1024，16，1）和（128，128，1）是否有区别？

两个都是开了16384个 block，可能从thread的角度来看，没什么区别。

我们以同样的尺寸，用tvm模仿一下cutlass，这里我选择grid size是（128，128，1），block size是（16，16，1）。

首先，切一下KTile，这里选择的大小是16，因为我们一个block需要load BK*BN + BM * BK 大小的数据，在这个尺寸下是16128，如果选择BK=8，则是8\128, tvm把里面那个thread直接bind到thread会导致有一半的thread不能参与load global memory的任务里来：

    BK = 16

    ## thread block tiling
    ko, ki = s[CC].split(k, factor=BK)
    xc, yc = s[CC].op.axis
    s[CC].reorder(ko, ki, xc, yc)

对应到生成的代码：

      for (k.outer: int32, 0, 2048) {
        for (k.inner: int32, 0, 8) {
          for (ii.c: int32, 0, 8) {
            for (jj.c: int32, 0, 8) {
              let cse_var_4: int32 = ((ii.c*8) + jj.c)
              let cse_var_3: int32 = ((k.outer*64) + (k.inner*8))
              C.local_1[cse_var_4] = (C.local_1[cse_var_4] + (A.shared.local_1[(cse_var_3 + jj.c)]*B.shared.local_1[(cse_var_3 + ii.c)]))
            }
          }
        }
      }

一行有16384个数据，由于BK是8，则每个thread需要迭代2048次才能玩车个完整的C的计算，很正确。

与之前的schedule不一样的是，一个thread需要同时处理多个元素，所以需要首先把loop拆分到threadIdx上：

    bx, xi = s[C].split(C.op.axis[0], nparts=Grid_Size_X)
    by, yi = s[C].split(C.op.axis[1], nparts=Grid_Size_Y)
    s[C].bind(bx, block_x)
    s[C].bind(by, block_y)
    s[C].reorder(bx, by, xi, yi)

接下来可以把A和B的cache也加上，这下不会out of shared memory了，理一下思路：

在 k.outer 这个维度的迭代中，每次从global memory load BM*BK*BN个数据load到share memory，显然这部分功能可以每个先线程load一部分加快速度。
在k.inner 这个维度的迭代中，把每个thread需要的部分数据load到local（register）上。

所以，我们先把两部分的load放到对应的位置:

    s[AA].compute_at(s[CC], ko)
    s[BB].compute_at(s[CC], ko)
    s[AL].compute_at(s[CC], ki)
    s[BL].compute_at(s[CC], ki)

对应到生成的代码：

for (k.outer: int32, 0, 1024) {
      for (ax0: int32, 0, 16) {
        for (ax1: int32, 0, 128) {
          A.shared_1: Buffer(A.shared, float32, [2048], [], scope="shared")[((ax0*128) + ax1)] = A[((((k.outer*262144) + (ax0*16384)) + (blockIdx.y*128)) + ax1)]
        }
      }
      for (ax0_1: int32, 0, 16) {
        for (ax1_1: int32, 0, 128) {
          B.shared_1: Buffer(B.shared, float32, [2048], [], scope="shared")[((ax0_1*128) + ax1_1)] = B[((((k.outer*262144) + (ax0_1*16384)) + (blockIdx.x*128)) + ax1_1)]
        }
      }
      for (k.inner: int32, 0, 16) {
        for (ax1_2: int32, 0, 8) {
          A.shared.local_1: Buffer(A.shared.local, float32, [8], [], scope="local", align=32)[ax1_2] = A.shared_1[(((k.inner*128) + (threadIdx.y*8)) + ax1_2)]
        }
        for (ax1_3: int32, 0, 8) {
          B.shared.local_1: Buffer(B.shared.local, float32, [8], [], scope="local", align=32)[ax1_3] = B.shared_1[(((k.inner*128) + (threadIdx.x*8)) + ax1_3)]
        }
        for (ii.c: int32, 0, 8) {
          for (jj.c: int32, 0, 8) {
            let cse_var_1: int32 = ((ii.c*8) + jj.c)
            C.local_1[cse_var_1] = (C.local_1[cse_var_1] + (A.shared.local_1[jj.c]*B.shared.local_1[ii.c]))
          }
        }
      }
    }

直接生成cuda kernel运行，一次迭代消耗的时间是13717.8 ms，是cutlass的30倍。

然后，完成读取global memory的并行加速：

    aa_tx, aa_xi = s[AA].split(s[AA].op.axis[0], nparts=Block_Size_X)
    aa_ty, aa_yi = s[AA].split(s[AA].op.axis[1], nparts=Block_Size_Y)
    s[AA].reorder(aa_tx, aa_ty, aa_xi, aa_yi)
    s[AA].bind(aa_tx, thread_x)
    s[AA].bind(aa_ty, thread_y)

    bb_tx, bb_xi = s[BB].split(s[BB].op.axis[0], nparts=Block_Size_X)
    bb_ty, bb_yi = s[BB].split(s[BB].op.axis[1], nparts=Block_Size_Y)
    s[BB].reorder(bb_tx, bb_ty, bb_xi, bb_yi)
    s[BB].bind(bb_tx, thread_x)
    s[BB].bind(bb_ty, thread_y)

生成的代码：

for (k.outer: int32, 0, 1024) {
      for (ax1.inner: int32, 0, 8) {
        A.shared_1: Buffer(A.shared, float32, [2048], [], scope="shared")[(((threadIdx.x*128) + (threadIdx.y*8)) + ax1.inner)] = A[(((((k.outer*262144) + (threadIdx.x*16384)) + (blockIdx.y*128)) + (threadIdx.y*8)) + ax1.inner)]
      }
      for (ax1.inner_1: int32, 0, 8) {
        B.shared_1: Buffer(B.shared, float32, [2048], [], scope="shared")[(((threadIdx.x*128) + (threadIdx.y*8)) + ax1.inner_1)] = B[(((((k.outer*262144) + (threadIdx.x*16384)) + (blockIdx.x*128)) + (threadIdx.y*8)) + ax1.inner_1)]
      }
      for (k.inner: int32, 0, 16) {
        for (ax1: int32, 0, 8) {
          A.shared.local_1: Buffer(A.shared.local, float32, [8], [], scope="local", align=32)[ax1] = A.shared_1[(((k.inner*128) + (threadIdx.y*8)) + ax1)]
        }
        for (ax1_1: int32, 0, 8) {
          B.shared.local_1: Buffer(B.shared.local, float32, [8], [], scope="local", align=32)[ax1_1] = B.shared_1[(((k.inner*128) + (threadIdx.x*8)) + ax1_1)]
        }
        for (ii.c: int32, 0, 8) {
          for (jj.c: int32, 0, 8) {
            let cse_var_1: int32 = ((ii.c*8) + jj.c)
            C.local_1[cse_var_1] = (C.local_1[cse_var_1] + (A.shared.local_1[jj.c]*B.shared.local_1[ii.c]))
          }
        }
      }
    }

一次迭代的时间缩小到854ms，已经缩小到cutlass的两倍了！

Wrap Tile & Bank Conflict

Shared memory 是以 4 bytes (一个word) 为单位分成 banks：

每个bank的带宽是32 bits per clock cycle
连续的32个word是放在连续的32个banks中

有关bank conflict的内容网上有很多说明的文章，下面是参考这一篇给出的说明，3090上，一共有32个bank，下面是以16个bank来举例的：

因此，假设以下的数据：

__shared__ int data[128];

那么，data[0] 是 bank 0、data[1] 是 bank 1、data[2] 是 bank 2、…、data[15] 是 bank 15，而 data[16] 又回到 bank 0。

由于 warp 在执行时是以 half-warp 的方式执行（关于half-warp参照上一篇博客），因此分属于不同的 half warp 的 threads，不会造成 bank conflict。

因此，如果程序在存取 shared memory的时候，使用以下的方式：

int number = data[base + tid];

那就不会有任何 bank conflict，可以达到最高的效率。但是，如果是以下的方式：

int number = data[base + 4 * tid];

那么，thread 0 和 thread 4 就会存取到同一个 bank，thread 1 和 thread 5 也是同样，这样就会造成 bank conflict。在这个例子中，一个 half warp 的 16 个 threads 会有四个 threads 存取同一个 bank，因此存取 share memory 的速度会变成原来的 1/4。

一个重要的例外是，当多个 thread 存取到同一个 shared memory 的地址时，shared memory 可以将这个地址的 32 bits 数据「广播」到所有读取的 threads，因此不会造成 bank conflict。

例如：

int number = data[3];

这样不会造成 bank conflict，因为所有的 thread 都读取同一个地址的数据。

很多时候 shared memory 的 bank conflict 可以透过修改数据存放的方式来解决。

抱着这个思路回过头来观察我们现在的kernel对share memory的读写：

    # write
    for (ax1.inner: int32, 0, 8) {
        A.shared_1: Buffer(A.shared, float32, [2048], [], scope="shared")[(((threadIdx.x*128) + (threadIdx.y*8)) + ax1.inner)] = A[(((((k.outer*262144) + (threadIdx.x*16384)) + (blockIdx.y*128)) + (threadIdx.y*8)) + ax1.inner)]
    }
    for (ax1.inner_1: int32, 0, 8) {
        B.shared_1: Buffer(B.shared, float32, [2048], [], scope="shared")[(((threadIdx.x*128) + (threadIdx.y*8)) + ax1.inner_1)] = B[(((((k.outer*262144) + (threadIdx.x*16384)) + (blockIdx.x*128)) + (threadIdx.y*8)) + ax1.inner_1)]
    }
    # read
    for (ax1: int32, 0, 8) {
        A.shared.local_1: Buffer(A.shared.local, float32, [8], [], scope="local", align=32)[ax1] = A.shared_1[(((k.inner*128) + (threadIdx.y*8)) + ax1)]
    }
    for (ax1_1: int32, 0, 8) {
        B.shared.local_1: Buffer(B.shared.local, float32, [8], [], scope="local", align=32)[ax1_1] = B.shared_1[(((k.inner*128) + (threadIdx.x*8)) + ax1_1)]
    }

因为float32的长度刚好是一个word，则索引%32就可以作为bank下标的标识。

写的过程中，假设第0个block的第0个thread开始，写入的地址为 0, 8, 16, 24, 32, 40…,明显第一个thread和第五个thread就写入了同一块bank。

thread虽然在逻辑上都是并行处理，但实际上是调度到warp上快速执行完一组thread，再快速执行下一组，而且bank是simd级别的，可以理解为一个share memory controller，所以我们只需要保证同一组warp上的thread不conflict就可以了。

分析之前的程序，写share memory部分造成conflict的原因主要是写schedule的时候，没有把最后一层的loop拆成threadx导致的，重新改一下schedule：

    aa_yi, aa_ty = s[AA].split(s[AA].op.axis[0], factor=Block_Size_Y)
    aa_xi, aa_tx = s[AA].split(s[AA].op.axis[1], factor=Block_Size_X)
    s[AA].reorder(aa_ty, aa_tx, aa_xi, aa_yi)
    s[AA].bind(aa_ty, thread_y)
    s[AA].bind(aa_tx, thread_x)

    bb_xi, bb_tx  = s[BB].split(s[BB].op.axis[0], factor=Block_Size_Y)
    bb_yi, bb_ty  = s[BB].split(s[BB].op.axis[1], factor=Block_Size_X)
    s[BB].reorder(bb_ty, bb_tx, bb_xi, bb_yi)
    s[BB].bind(bb_ty, thread_y)
    s[BB].bind(bb_tx, thread_x)

生成的代码：

    for (int ax1_outer = 0; ax1_outer < 8; ++ax1_outer) {
      A_shared[(((((int)threadIdx.y) * 128) + (ax1_outer * 16)) + ((int)threadIdx.x))] = A[(((((k_outer * 262144) + (((int)threadIdx.y) * 16384)) + (((int)blockIdx.y) * 128)) + (ax1_outer * 16)) + ((int)threadIdx.x))];
    }
    for (int ax1_outer1 = 0; ax1_outer1 < 8; ++ax1_outer1) {
      B_shared[(((((int)threadIdx.x) * 128) + (ax1_outer1 * 16)) + ((int)threadIdx.y))] = B[(((((k_outer * 262144) + (((int)threadIdx.x) * 16384)) + (((int)blockIdx.x) * 128)) + (ax1_outer1 * 16)) + ((int)threadIdx.y))];
    }

可以发现变成 s[tid] = g[xx]的形式了，没有bank conflict，而运行的时间也快了一些(快了100ms 左右的样子)。

显然读的过程也存在bank conflict，这一点可以从第三张图中看到，左上角的32个虚线方格是一个warp的32个thread，同一行的thread会同时读取同一个A的值，同一列的thread会同时读取一个B的值，而从图中也能看出cutlass是怎么缓解这个问题：

访存的地址是：

      for (int ax1 = 0; ax1 < 8; ++ax1) {
        A_shared_local[ax1] = A_shared[(((k_inner * 128) + (((int)threadIdx.y) * 8)) + ax1)];
      }
      for (int ax11 = 0; ax11 < 8; ++ax11) {
        B_shared_local[ax11] = B_shared[(((k_inner * 128) + (((int)threadIdx.x) * 8)) + ax11)];
      }

对于A来说，前十六个thread访问的memory地址是一样的，但是这时候应该有broadcast，对于B来说，访问的地址是0，8，16，32（32的时候会有conflict），

之前有人已经给出更丰富的图片说明了：https://zhuanlan.zhihu.com/p/518857175

在tvm里，实现这样的切分方式，并且明确说明了目的是为了缓解bank conflict的是virtual thread，如下：

    tyz, yi = s[C].split(yi, nparts=2)
    ty, yi = s[C].split(yi, nparts=Block_Size_Y)
    txz, xi = s[C].split(xi, nparts=8)
    tx, xi = s[C].split(xi, nparts=Block_Size_X)
    s[C].bind(tyz, thread_yz)
    s[C].bind(txz, thread_xz)
    s[C].bind(ty, thread_y)
    s[C].bind(tx, thread_x)

565ms，生成的代码是：

        B.shared.local_1: Buffer(B.shared.local, float32, [1], [], scope="local", align=4)[0] = B.shared_1[((k.inner*128) + threadIdx.x)]
        B.shared.local_1[1] = B.shared_1[(((k.inner*128) + threadIdx.x) + 16)]
        B.shared.local_1[2] = B.shared_1[(((k.inner*128) + threadIdx.x) + 32)]
        B.shared.local_1[3] = B.shared_1[(((k.inner*128) + threadIdx.x) + 48)]
        B.shared.local_1[4] = B.shared_1[(((k.inner*128) + threadIdx.x) + 64)]
        B.shared.local_1[5] = B.shared_1[(((k.inner*128) + threadIdx.x) + 80)]
        B.shared.local_1[6] = B.shared_1[(((k.inner*128) + threadIdx.x) + 96)]
        B.shared.local_1[7] = B.shared_1[(((k.inner*128) + threadIdx.x) + 112)]

这样，就没有bank conflict了，vthread虽然是虚拟的，但是如果直接用split+reorder的方式来替换：

    tyz, yi = s[C].split(yi, nparts=2)
    ty, yi = s[C].split(yi, nparts=Block_Size_Y)
    txz, xi = s[C].split(xi, nparts=2)
    tx, xi = s[C].split(xi, nparts=Block_Size_X)
    s[C].bind(ty, thread_y)
    s[C].bind(tx, thread_x)
    s[C].reorder(tyz, txz, ty, tx, xi, yi)

速度反而会慢不少（692.994ms），也有人在论坛里讨论过,而且看tvm的论文里提到的vthread，似乎还有类似于分析依赖，软件流水的神奇功能。

其他

vectorize

从global memory读取数据可以使用lgd.128指令，一次读4个float32的数据，从share memory 读取数据，可以用lgs.128，对应到tvm里则是使用vectorize进行向量化，这样，一条访存指令就可以直接load 4个float32的数据，为了做到这一点，首先需要从循环中把可以vectorize的shape手动拆出来，再进行向量化：

    aa_yi, aa_ty = s[AA].split(s[AA].op.axis[0], factor=Block_Size_Y)
    aa_xi, aa_tx = s[AA].split(s[AA].op.axis[1], factor=Block_Size_X * 4)
    aa_tx, aa_vi = s[AA].split(aa_tx, nparts=Block_Size_X)
    s[AA].reorder(aa_ty, aa_tx, aa_yi, aa_xi, aa_vi)
    s[AA].bind(aa_ty, thread_y)
    s[AA].bind(aa_tx, thread_x)
    s[AA].vectorize(aa_vi)
    bb_yi, bb_ty = s[BB].split(s[BB].op.axis[0], factor=Block_Size_Y)
    bb_xi, bb_tx = s[BB].split(s[BB].op.axis[1], factor=Block_Size_X * 4)
    bb_tx, bb_vi = s[BB].split(bb_tx, nparts=Block_Size_X)
    s[BB].reorder(bb_ty, bb_tx, bb_yi, bb_xi, bb_vi)
    s[BB].bind(bb_ty, thread_y)
    s[BB].bind(bb_tx, thread_x)
    s[BB].vectorize(bb_vi)

生成的代码：

    for (int ax1_outer = 0; ax1_outer < 2; ++ax1_outer) {
      *(float4*)(A_shared + (((((int)threadIdx.y) * 128) + (ax1_outer * 64)) + (((int)threadIdx.x) * 4))) = *(float4*)(A + (((((k_outer * 262144) + (((int)threadIdx.y) * 16384)) + (((int)blockIdx.y) * 128)) + (ax1_outer * 64)) + (((int)threadIdx.x) * 4)));
    }
    for (int ax1_outer1 = 0; ax1_outer1 < 2; ++ax1_outer1) {
      *(float4*)(B_shared + (((((int)threadIdx.y) * 128) + (ax1_outer1 * 64)) + (((int)threadIdx.x) * 4))) = *(float4*)(B + (((((k_outer * 262144) + (((int)threadIdx.y) * 16384)) + (((int)blockIdx.x) * 128)) + (ax1_outer1 * 64)) + (((int)threadIdx.x) * 4)));
    }

514ms，又快了50ms，接下来还可以对local register的访问做vectorize，也可以用lds.128，所以需要把vthread的大小改成2x2，这样分出来刚好是两个可以被vectorize的向量：

    tyz, yi = s[C].split(yi, nparts=2)
    ty, yi = s[C].split(yi, nparts=Block_Size_Y)
    txz, xi = s[C].split(xi, nparts=2)
    tx, xi = s[C].split(xi, nparts=Block_Size_X)

    al_yi, al_xi = s[AL].op.axis
    s[AL].vectorize(al_xi)
    bl_yi, bl_xi = s[BL].op.axis
    s[BL].vectorize(bl_xi)

生成的代码：

__shared__ float4 A_shared[512];
__shared__ float4 B_shared[512];     
    for (int k_inner = 0; k_inner < 16; ++k_inner) {
      *(float4*)(A_shared_local + 0) = A_shared[((k_inner * 32) + ((int)threadIdx.y))];
      *(float4*)(A_shared_local + 4) = A_shared[(((k_inner * 32) + ((int)threadIdx.y)) + 16)];
      *(float4*)(B_shared_local + 0) = B_shared[((k_inner * 32) + ((int)threadIdx.x))];
      *(float4*)(B_shared_local + 4) = B_shared[(((k_inner * 32) + ((int)threadIdx.x)) + 16)];
      for (int ii_c = 0; ii_c < 4; ++ii_c) {
        for (int jj_c = 0; jj_c < 4; ++jj_c) {
          C_local[((ii_c * 4) + jj_c)] = (C_local[((ii_c * 4) + jj_c)] + (A_shared_local[jj_c] * B_shared_local[ii_c]));
          C_local[(((ii_c * 4) + jj_c) + 32)] = (C_local[(((ii_c * 4) + jj_c) + 32)] + (A_shared_local[(jj_c + 4)] * B_shared_local[ii_c]));
          C_local[(((ii_c * 4) + jj_c) + 16)] = (C_local[(((ii_c * 4) + jj_c) + 16)] + (A_shared_local[jj_c] * B_shared_local[(ii_c + 4)]));
          C_local[(((ii_c * 4) + jj_c) + 48)] = (C_local[(((ii_c * 4) + jj_c) + 48)] + (A_shared_local[(jj_c + 4)] * B_shared_local[(ii_c + 4)]));
        }
      }
    }

这样，一个warp里的thread访问的bank就是0，4，8，12，16…32. 到第九个线程的时候还是会出现bank conflict？这个应该是无法避免了，可能得交给nvcc来做软件流水，但是速度确实快了不少，达到了423ms, cublas和cutlass测出来的结果都是420ms,已经基本上达到这个水平了！

double buffer

当然如果load的时间过长，计算不能hide load latency的时候，可以试试double buffer来进行乒乓缓存，不过占用的share memory会翻倍，直接：

    s[AA].double_buffer()
    s[BB].double_buffer()

但是在我的case下面，这个操作加上反而变得更慢（430ms）。。

最后

都已经能打到和cutlass相同的水平，我觉得也就差不多了，至少验证了tvm是真的行（至少在gemm上），当然以上有一些我自己的思考，不一定对，文章里还有一些问题没有解决，比如，（1024，16，1）和（128，128，1）这两个grid size是一样的么？群友觉得是不一样的，而我图编程方便，使用的后者，cutlass用的是前者，但从结果上来看，效果还可以，关于这些，欢迎大佬指正与交流，至于前面的区域，以后有时间了再来探索吧。

#CUDA Programming #MLSys

a hgemm tvm schedule

naive gemm

Blocked Gemm

Thread Block Tile

Wrap Tile & Bank Conflict

其他

vectorize

double buffer

最后

Comments

Your browser is out-of-date!