9. MUSA性能优化

在本章中，我们将讲解一些 MUSA 中常用的计算优化和访存优化的方法，同时会介绍性能调优的相关方法和工具。最后，我们还会提供 Reduction、GEMM 等优化实例供参考和学习。

9.1. 核函数优化

9.1.1. 并行度优化

9.1.1.1. 最大化计算并行度

为了最大化计算并行度，首先我们需要选择合适的算法。对于给定的应用场景，不同算法的计算并行度可能会有很大差异。有些天然适合并行计算的场景，例如 reduce；有些场景的常规算法则不适合进行并行实现，例如 sort，这时我们则需要探索其适合并行实现的算法，比如可以使用 双调排序 Bitonic sort 来进行并行实现；对于一些无法直接进行并行实现的复杂场景，我们可以对其进行拆分，选取部分功能进行并行实现，以提升整体的性能。

9.1.1.2. 减少分支

在一个 thread block 中，由于 if、switch、for 和 while 等分支控制语句的使用，会导致不同线程的执行路径产生 divergence，从而导致不同线程在某些时刻需要执行不同的指令。这种情况会导致线程之间产生额外的等待开销，影响程序的性能。

• 循环展开

  • 采用循环展开可以减少循环开销和分支，从而提高性能。循环展开可以手动进行，也可以使用编译器优化来自动完成。循环展开通常可以提高计算性能，但需要注意循环展开后指令的条数，指令条数过多容易造成 Instruction Cache Misses，反而会降低核函数的性能。

  #pragma unroll
      for (int i = 0; i < 10; ++i)
      {
          // ...
      }

• 我们还可以使用一些技巧来避免分支，例如使用条件运算符（ternary operator）代替 if-else 语句。

9.1.1.3. 向量化数据读取

向量化是一种针对数据并行的优化技巧。通过向量化数据读取，每个线程可以同时读取多个相邻的数据元素，从而减少读取操作和访存延迟，提高程序的效率。

9.1.2. 核函数执行配置优化

MUSA 定义的线程结构分为三级：grid、block 和 thread，采用这种层次化的线程结构是为了能有效的与 GPU 层次化的硬件结构相对应。在进行核函数的实现时，合理的核函数配置能最大限度地发挥硬件的性能。

9.1.2.1. 最大化硬件占用率

为了获取更好的性能，最大化硬件利用率是关键之一。在核函数的执行过程中，大量的 warp 可以提供高度的并行性，使得处理器可以借助 warp 切换来进行访存等延时隐藏。理论上warp数越多，并行度会越高，对于延时的隐藏效果会更好。但在实际应用中，并行度会受硬件资源总数、每个线程所需的资源数量（寄存器等）以及每个block所需的资源数量（shared memory等）决定的。例如，在 S80 显卡中，一个 MP 提供了 28KB 的 shared memory，若单个 block 使用 4KB 的 shared memory，在仅考虑 shared memory 的约束下，则最多可以同时运行 7 个 block。

active warps 个数受限于木桶理论，即受制于各个资源的最短板。因此在核函数的编写中，需要综合考虑单个线程的寄存器使用个数、单个block的shared memeory使用大小等，确保 active warps 的个数在一个合理的区间，以保证运行时有足够的 warp 数量以隐藏延时，最大化硬件占用率。

9.1.2.2. block size配置策略

block size 应设置为 warp size 的整数倍（SUDI 和 QY 架构中，warp size 为 128），因为如果 block size 不是 warp size 的整数倍，那么在执行时会导致一些 warp 只有部分线程被使用，从而浪费计算资源。将 block size 设置为 warp size 的整数倍可以最大化核函数的并行度以提升效率。

• 一般情况下，128、256、512 和 1024 可以作为候选的 block size 以获取较优的性能。
• GPU 一般会包含多个 MP，因此我们在进行核函数配置时，一半还需要考虑 block 的总个数，至少保证每个 MP 有一个 block 去执行。理想的状态是，active warps 的数量足够多，使得 GPU 可以通过 warp 切换隐藏延时。

如我们上面描述的，选择合适的 block size 有许多需要考虑的因素，这个过程需要编程人员具有一定的 GPU 编程经验。幸而还有一种适合新手的方法——auto fine-tune block size。具体来说，对于给定的核函数，该方法通过试验不同的 block size 大小，并利用如吞吐量、延时等性能监测指标来评估性能，从而确定较优的 block size 大小。很多开源的项目中都集成了 atuo fine-tune block size 的功能，例如 MNN、TNN 以及 ROCm。

需要注意的是，如果单个 block 所需的寄存器或 shared memory 超过单个 MP 提供的最大范围时，核函数会启动失败。

9.1.2.3. 多个核函数并发执行

在某些应用场景下，我们可以将无依赖关系的多个核函数分发至不同的 stream，使多个独立的核函数可以同时执行，以此来提高硬件利用率。

9.1.3. 空间换时间

9.1.3.1. double buffer

double buffer技术是一种通过交替使用两个缓冲区来实现数据传输和计算并行的技术。

• 在GPU计算中，通常需要进行数据传输和计算两个步骤，double buffer技术可以将这两个步骤并行执行，从而提高计算效率。具体来说，double buffer技术需要使用两个缓冲区来存储数据，一个缓冲区用于计算，另一个缓冲区用于数据传输。在计算时，使用其中一个缓冲区进行计算，同时在另一个缓冲区进行数据传输。计算完成后，再切换缓冲区，继续进行计算和数据传输。使用double buffer技术的优点是可以在数据传输和计算之间实现并行，提高计算效率。此外，double buffer 还可以减少 GPU 的空闲时间，提高 GPU 的利用率。
• double buffer 的缺点是需要额外的存储空间，我们需要平衡存储空间和计算效率的关系。

9.1.3.2. img2col

img2col是采用矩阵乘来实现卷积计算的步骤之一。具体做法为：

• 将4维的 feature map 和 weight 提前进行 img2col 转化为2维的矩阵，并使用额外的空间存储 img2col 的结果，这样在矩阵乘阶段我们可以直接使用。img2col 带来的好处是提高了矩阵乘阶段的数据访问局部性，降低了数据访问延迟，从而提高计算效率。当然，img2col 会引入一定的计算（weight 的 img2col 可以提前进行，而 feature map 的 img2col 则只能在推理阶段进行）和内存开销，我们需要平衡计算和内存消耗的开销。

9.1.4. 精度选择

GPU 对不同精度的数据具有不同的处理速度。针对不同的应用场景，我们可以在可接受的计算精度范围内选择较低的数据精度以提升应用的性能。

• 对于精度要求不高的应用可以通过低精度浮点数来提高计算性能。例如，使用单精度浮点数代替双精度浮点数，或者半精度（half）或混合精度（mixed precision）代替单精度浮点数。
• 使用量化模型提升推理的性能。相对于浮点数，定点数的吞吐更高，因此在可接受的模型精度丢失情况下，采用量化模型可以带来可观的性能提升。并且量化模型还具有轻量，低功耗等优点。

9.1.5. 特殊计算单元

• Tensor Core
  • Tensor Core是一种针对矩阵计算优化的硬件加速器。使用 Tensor Core 可以大幅提高矩阵乘法的计算性能，但需要注意数据类型和格式的兼容性。在QY系列显卡中，我们提供了对mma的支持，例如16x8x16 IMMA指令。

9.2. 访存优化

MUSA可用的存储器的速度由快到慢是, 通用寄存器 > 常量内存 > 共享内存 > L1/L2/LLC > 全局内存. 优化GPU访存的核心是让操作尽量发生在速度快的存储器中. 注意, 共享内存和全局内存的访问都存在访存合并行为, 即连续的线程访问连续的地址则请求可以合并, 以充分利用存储器的带宽. 所有线程访问相同的地址会触发广播, 即实际只有一条真实的访存请求被下发给存储器, 其请求的数据会被分发给不同的线程, 降低存储器的压力.

• reduce变量尽量在寄存器上reduce完再写回全局内存
• block内不同线程reduce数据通过共享内存进行
• 在共享内存中做转置保证不同线程读写全局内存可以合并, 一般情况下共享内存的速度远快于全局内存
• 访存和计算是不同的硬件单元, 可以通过double buffer等手段在线程层面让访存和计算并行, 也可以通过提高occupancy增加线程间的计算和访存并行
• 当线程间对全局内存的访问无法合并时应尽量提升单个线程访问全局内存的量, 以提升cacheline利用率

9.3. inline MUSA IR

> 敬请期待后续版本更新

9.4. 性能调优方法及工具

9.4.1. mSight

> 敬请期待后续版本更新

9.4.2. 自定义性能统计与分析

9.4.2.1. Kernel Level Roofline

首先Kernel Level Roofline Model主要描述了某个kernel在一个计算平台的限制下，到底能达到多快的浮点计算速度。更具体的来说，它解决的，是“计算量为A且访存量为B的kernel在算力为C且带宽为D的计算平台所能达到的理论性能上限E是多少”。

1. 计算平台指标

   • 算力$\pi$：指的是一个计算平台倾尽全力每秒钟所能完成的浮点运算数。单位是 $FLOPS$ or $FLOP/s$。
   • 带宽$\beta$：指的是一个计算平台倾尽全力每秒所能完成的内存交换量。单位是$Byte/s$。
   • 计算强度上限$I_{max}$：它描述的是在这个计算平台上，单位内存交换最多用来进行多少次计算。单位是$FLOPS/Bytes$:
     $$I_{max} = {\pi \over \beta }$$

2. Kernel指标

   • 计算量：指的是执行某个kernel所发生的浮点运算个数。单位是 $FLOP$或者$FLOPs$。
   • 访存量：指的是输入单个样本，执行kernel所发生的内存交换总量。
   • 计算强度: 由计算量除以访存量就可以得到kernel的计算强度，它表示此kernel在计算过程中，每Byte内存交换到底用于进行多少次浮点运算。单位是$FLOPs/Byte$。

3. Roofline Model

   Roofline Model如图1所示，我们可以把Roofline Model分为两个区域:

   • 计算瓶颈区：当模型的计算强度超过平台的理论最大计算强度的时候，算力最多只能达到计算平台的算力;反之如果计算密度较大，程序性能受硬件最大计算峰值限制，称为计算密集型程序，即图中绿色区域。此时性能上界=硬件算力，表现为图中的横线。此时计算速度不受计算密度影响，但计算密度越大，所需内存带宽就越少。
   • 带宽瓶颈区表示：当模型的计算强度小于平台的理论最大计算强度的时候，算力是受平台带宽限制的（假设计算强度固定）。当程序的计算密度较小时，程序访存多而计算少，性能受内存带宽限制，称为访存密集型程序，即图中红色色区域。在此区域的程序性能上界=计算密度×内存带宽，表现为图中的斜线，其中斜率为内存带宽的大小。计算密度越大，程序所能达到的速度上界越高，但使用的内存带宽始终为最大值。
     

图1

4. 基于Roofline的Kernel性能分析

   在真实世界中，核函数都必须依赖于具体的计算平台才能展现自己的实力。这样他们和计算平台的"默契程度"才能决定核函数的实际表现，而Kernel Level Roofline恰好反映了Kernel和计算平台之间的关系。根据计算平台的理论性能、kernel 的计算强度，和实际运行中 kernel 的计算效率就可以画出
   kernel level roofline model。如图2, 其中横轴代表 kernel 的计算强度，而纵轴就是该 kernel 所能达到的最大计算性能。我们可以把图像分为以下几个区域：

   • 粉色线以上和蓝色线以左的区域是程序无论如何都无法达到的性能，因为它意味着超过了计算机的峰值计算性能/访存带宽。
   • 淡蓝色区域（计算密度小于$I_{max}$点）是性能较好的访存密集型程序，这部分程序的访存带宽利用率较高，评价访存密集型程序的指标主要选用访存带宽。
   • 淡粉色区域（计算密度大于$I_{max}$点）是性能较好的计算密集型程序，有较好的数据重用率与数据局部性，这部分程序的浮点性能较高，评价计算密集型程序的指标主要选用浮点性能。
   • 红色虚线部分，带宽与浮点性能都远低于峰值性能的kernels，如果程序性能处于这个部分，需要考虑优化算法提高性能，达到粉色或者蓝色区域。

   至于那些贴近于roofline的kernel较好的利用了计算资源（绿点），而那些远离roofline的kernel（红点）则是我们重点需要去优化、提高计算资源利用率的。这里需要注意，对比FLOP/s很低的kernel（例如第一个绿点），红点虽然看起来FLOP/s更高，但是比绿点更有优化性价比。红点离roofline还有较远距离，可通过不断优化访存函数、计算函数来提高访存带宽/浮点计算性能利用率（可根据坐落的区域判断）。除此之外，如果想要提高FLOP/s很低但是又比较接近roofline的kernel，可以通过改进计算方法/减少数据传输时间来提高计算密度（例如提高空间局部性、提高cache命中率、改进数据结构、数据类型）。
   

图2

上述的方案中，我们默认带宽和访存量都是针对的GPU device memory，但其实kernel的性能是受以下四种memory限制的：

• Registers
• L1，L2，cache
• HBM (GPU device memory)
• DDR (host memory)
  

图3

不同的memory具有不同的带宽，而同一个kernel在不同的memory level下有不同的访存量。图3就展示了不同memory level下的roofline情况帮助我们分析kernel在不同存储结构下的性能表现。首先图中画了三条不同memory level下的roofline。同一个kernel是用同一种颜色标识的，从图中可以发现同一个kernel不管是在哪个memory level下具有相同的计算性能（因为在同一个计算平台下）, 但由于不同的访存量所以具有不同的计算密度。同一个kernel在不同memory level下之间的距离可以衡量cache的利用率，如果同一个kernel在不同level下之间距离很短，说明cache利用率不高，如果很长，说明cache利用率高。除此之外，也可以单独分析每个kernel在某个存储结构下的性能瓶颈，从而提高kernel性能。

图4

5. 基于指令级别Ceilings的kernel Roofline分析

   Roofline model展示了性能的上界，那如果性能表现远低于roofline的话，我们应该按什么顺序优化我们的程序呢?除了上述方案中对kernel程序本身做出优化，这里也列出了一些指令级别的performance ceilings，如果没有做到某个ceiling相关的优化，是不可能达到该ceiling的上界的。这里将ceiling分为计算瓶颈ceiling和带宽瓶颈ceiling。

   • 计算瓶颈的优化：
   1. 优化ILP( instruction level parallelism)和应用SIMD：
      • 指令级并行（ ILP, Instruction Level Parallelism）是指利用流水级并行和多指令发射等方式提高程序执行的并行度；
      • 数据级并行（DLP, Data Level Parallelism）是指处理器能够同时处理多条数据的并行方式，即SIMD。
   2. floating-point balanced（平衡浮点运算组合）：要达到浮点运算性能表现，需要乘法和加法的次数尽可能一样，因为许多的计算平台有专门的multiply-add指令。
   • 带宽瓶颈优化
   1. Restructure loops for unit stride accesses（重构单位步幅访问循环）
   2. Ensure memory affinity（确保内存亲和性）：这个 优化分配数据和分配给该数据的线程到相同的内存处理器对，以便处理器很少需要访问连接到其他芯片的内存。
   3. Use software prefetching（使用软件预取）：在某些计算机上， 软件预取提供比硬件单独预取带来更多的带宽。

   图5（a）展示了计算瓶颈ceilings，说明了在imbalance floating-point下， 算力最多达到8.8GFlops/s, 而在没有优化ILP和SIMD下，算力最多只能达到2.2GFlops/s。（b）展示了带宽瓶颈ceilings，道理同上。而（C）是两种ceilings的混合，他告诉我们在不同的计算强度下，该如何选择优化方式。首先对某个kernel计算其计算强度，再对该点做垂线，与该线能交汇的ceilings则是可以优化的方向。比如kernel2只和计算瓶颈ceilings相交，所以它只能做计算瓶颈的指令优化。而kernel1，既可做计算瓶颈优化，也可做么带宽瓶颈优化。
   

图5

9.4.3. mcc性能调优支持

mcc提供了包括优化选项、函数级别的属性、代码级别的注解等方式，为MUSA代码性能调优提供支持。

9.4.3.1. 优化选项

优化选项一般由-mllvm开头，每个优化选项开关前都需要加一个-mllvm作为前置修饰。

`-mtgpu-maxregcnt=

`

为本次编译设置最大temporary register使用数量限制。

使用示例：

mcc axpy.mu -lmusart -L/usr/local/musa/lib -mllvm -mtgpu-maxregcnt=256

9.4.3.2. 函数属性

mtgpu-num-usreg

该属性修饰kernel函数，用于限制该kernel使用的temporary register数量。“0”表示对temporary register数量不进行限制。

使用示例：

__global__ attribute((mtgpu-num-usreg(256))) void axpy(float *a, float *b, float c) {
    a[threadIdx.x] = b[threadIdx.x] * c;
}

mtgpu-unroll_threashold

该属性修饰kernel函数，用于限制kernel中的循环展开阈值。使用方式与mtgpu-num-usreg相同。

9.5. 优化实例

9.5.1. 优化实例1 - Reduction

9.5.1.1. 使用全局内存

并行归约（Reduction）是一种基础的并行算法，简单来说，我们有N个输入数据，使用一个符合结合律的二元操作符作用其上，最终生成1个结果。这个二元操作符可以是求和、取最大、取最小、平方、逻辑与或等等。
由于加法的交换律和结合律，数组可以以任意顺序求和。所以我们会自然而然产生这样的思路：首先把输入数组划分为更小的数据块，之后用一个线程计算一个数据块的部分和，最后把所有部分和再求和得出最终结果。

假设数组元素总数为2的整数次方，我们就可以将数组后半部分的各个元素与前半部分对应的数组相加。重复此过程，最后得到的第一个数就是最初数组的元素之和。这种方法叫做“折半归约”。如果使用一维网格和线程块，并将网格大小和线程块大小的乘积取N，可以写出一个简单的核函数，最后的结果就保存在 d_x[0]中：

void __global__ reduce(real *d_x, int N) {
    int n = blockDim.x * blockIdx.x + threadIdx.x;
    for (int offset = N / 2; offset > 0; offset /= 2) {
        if (n < offset) {
            d_x[n] += d_x[n + offset];
        }
    }
}

实际上该核函数并不能得到正确的结果，因为多线程的MUSA程序，两个不同线程中质量的执行顺序可能和代码中所展现的有所不同。

在循环的前两次迭代中有：

if (n < N / 2) {
    d_x[n] += d_x[n + n / 2];
}
 
if (n < N / 4) {
    d_x[n] += d_x[n + n / 4];
}

考虑第一次迭代中会有向数组元素 d_x[N / 4]写入数据的操作；在第二次迭代中会有从 d_x[N / 4]读取数据的操作，但是这两个迭代并不是在同一个线程中完成的，有一种可能的情况：当某线程执行第二行语句时，另一个执行第一行语句的线程还没有执行完毕。这种情况下，就有可能出现错误的结果。

那么如果要保证核函数中语句的执行顺序与出现顺序一致，则需要使用一种同步机制。MUSA中提供了一个同步函数 syncthreads()。这个函数只能用在核函数中，其常用方法就是不带任何参数的使用它。这个函数可以保证一个线程块中的所有线程，或者说是所有线程束在执行该语句后面的语句之前都执行了该语句前面的语句。但这个函数只能保证同一个线程块之内的线程同步，不同线程块之间的执行次序依旧是不确定的。基于这个函数我们实现正确的数组归约求和MUSA程序代码如下：

typedef float real;

const int NUM_REPEATS = 100;
const int N = 100000000;
const int M = sizeof(real) * N;
const int BLOCK_SIZE = 128;
 
void timing(real *h_x, real *d_x);
 
int main(void) {
    real *h_x = (real *)malloc(M);
    for (int n = 0; n < N; ++n) {
        h_x[n] = 1.01;
    }
    real *d_x;
    CHECK(musaMalloc(&d_x, M));
 
    printf("Using global memory only:\n");
    timing(h_x, d_x);
 
    free(h_x);
    CHECK(musaFree(d_x));
    return 0;
}
 
void __global__ reduce_global(real *d_x, real *d_y) {
    const int tid = threadIdx.x;
    real *x = d_x + blockDim.x * blockIdx.x;
 
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            x[tid] += x[tid + offset];
        }
        __syncthreads();
    }
 
    if (tid == 0) {
        d_y[blockIdx.x] = x[0];
    }
}
 
real reduce(real *d_x) {
    int grid_size = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    const int ymem = sizeof(real) * grid_size;
    const int smem = sizeof(real) * BLOCK_SIZE;
    real *d_y;
    CHECK(musaMalloc(&d_y, ymem));
    real *h_y = (real *)malloc(ymem);
 
    reduce_global<<<grid_size, BLOCK_SIZE>>>(d_x, d_y);
 
    CHECK(musaMemcpy(h_y, d_y, ymem, musaMemcpyDeviceToHost));
 
    real result = 0.0;
    for (int n = 0; n < grid_size; ++n) {
        result += h_y[n];
    }
 
    free(h_y);
    CHECK(musaFree(d_y));
    return result;
}
 
void timing(real *h_x, real *d_x) {
    real sum = 0;
 
    float total_time;
    for (int repeat = 0; repeat < NUM_REPEATS; ++repeat) {
        CHECK(musaMemcpy(d_x, h_x, M, musaMemcpyHostToDevice));
 
        musaEvent_t start, stop;
        CHECK(musaEventCreate(&start));
        CHECK(musaEventCreate(&stop));
        CHECK(musaEventRecord(start));
        musaEventQuery(start);
 
        sum = reduce(d_x);
 
        CHECK(musaEventRecord(stop));
        CHECK(musaEventSynchronize(stop));
        float elapsed_time;
        CHECK(musaEventElapsedTime(&elapsed_time, start, stop));
 
        if (repeat > 0) {
            total_time += elapsed_time;
        }
 
        CHECK(musaEventDestroy(start));
        CHECK(musaEventDestroy(stop));
    }
    printf("average time = %f.\n", total_time / NUM_REPEATS);
 
    printf("sum = %f.\n", sum);
}

这里仅对核函数进行相应解释：

核函数中定义了一个指针 x。赋值符号的右边是数组 d_x中第 blockDim.x * blockIdx.x个元素地址，如果了解指针的读者会理解其也可以写成 real *x = &d_x[blockDim.x * blockIdx.x]，实际上 x指向的就是全局内存中不同的地址，使得程序可以在不同的线程块中对数组中不同的地方进行归约，每个线程块处理 blockDim.x个数据。这里不再假设N是2的整数次方，但是N应该可以被 blockDim.x整除，且 blockDim.x为2的整数次方（代码中为128）。
接下来的循环语句就是归约算法的实现：在各个线程块内，对其中的数据独立的进行归约。其中 __syncthreads()保证同一线程块内的线程按照代码出现的顺序执行。
在这个核函数的操作中，将$10^8$的数组归约成长度为$10^8/128$的数组 d_y，那么为了计算整个数组之和，我们将数组 d_y从设备复制到主机，并在主机继续对其归约，得到最终结果。目前这个版本的实现不够高效，因为目前所有的工作并不是完全的在GPU上完成。我们暂时先采取这种方式。
在笔者的计算机上进行测试，在当前版本的MUSA编译环境下进行编译：

> mcc reduce_global_mem.mu -o reduce_global_mem -lmusart -O2

运行在MTT S80上使用单精度浮点数的运行时长为6.5ms，速度约为CPU计算的数倍。

9.5.1.2. 使用共享内存优化

上一小节的内容中的代码版本对全局内存的访问时很频繁的，在上一章的内容中提过，全局内存的访问速度是所有内存中最低的，实际应用中应该尽量减少对它的使用。所有设备内存里最高效的是寄存器，但是大量的数据导致寄存器的数量可能是不够的，尤其是需要线程和做的问题中应该使用对整个线程块可见的共享内存。

核函数中，若想定义一个共享内存变量，需要加上限定符 shared。在数组归约求和问题中，需要的是一个大小等于线程块大小的数组，于是我们可以定义变量：

__shared__ real s_y[BLOCK_SIZE];  // 变量名s前缀一般代表"shared"，同理d前缀意义为"device"，即设备上的变量，方便区分

若没有限定符修饰，则该变量可能是一段局部内存的变量。值得注意的是，在核函数中定义一个共享内存变量相当于在每个线程块中有一个该变量的副本，副本并不相同只是名称一样的变量，所有核函数中的操作都对各个线程块中的副本生效。

下面列出使用共享内存计算归约求和的核函数代码：

void __global__ reduce_shared(real *d_x, real *d_y)
{
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int n = bid * blockDim.x + tid;
    __shared__ real s_y[128];
    s_y[tid] = (n < N) ? d_x[n] : 0.0;
    __syncthreads();
 
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1)
    {
 
        if (tid < offset)
        {
            s_y[tid] += s_y[tid + offset];
        }
        __syncthreads();
    }
 
    if (tid == 0)
    {
        d_y[bid] = s_y[0];
    }
}

这里进行核函数的解析：

在核函数定义了共享内存数组 s_y，这里定义数组大小为线程块大小128。然后利用三目运算符进行判断，使其可以处理 N不是线程块大小整数倍的情况，并将对应的数组全局内存中的数据复制到共享内存中。这里注意 n的赋值，比如 bid取2时，会将全局内存中第 2 * blockDim.x个到第 3 * blockDim.x - 1个元素赋值给 s_y，即第2个线程块中的共享内存数组副本。
在赋值的下面利用同步函数 __syncthreads()进行线程块内的同步，确保共享内存变量中的数据对于归约操作之前全部准备就绪。
循环中对共享内存数组进行归约求和操作，替换了原来的全局变量。操作完成后，每个线程块中的 s_y[0]副本保存了对应的若干数组元素之和。
由于共享内存变量的生命周期只能维持在核函数中，所以必须在和函数执行完毕之前将数据保存到全局内存中。核函数最后的判断语句保证了在每个线程块中只被执行一次，操作的结果是将每个线程块中 s_y[0]保存到全局变量 d_y中对应的元素里，完整的代码中最后将 d_y中的元素转移到主机上然后再进行归约。
在这个核函数中我们使用了固定长度的共享内存数组，实际上MUSA也支持动态的共享内存。使用动态共享内存的好处就是在一些情况下时提高程序的可维护性，这里不做过多展开。将上述代码修改为动态共享数组只需要修改下面两处：

内核调用的写法：
在内核启动配置中增加第三个参数：

<<<grid_size, block_size, sizeof(real) * block_size>>> // 第三个参数为每个需要定义的动态共享内存的字节大小

注释中对第三个参数进行了解释。实际上在之前所有的内核启动配置中，这个参数都取默认值为0。

2. 核函数中对共享内存的声明：

extern __shared__ real s_y[];

主要是注意两点：必须加上限定词 extern，而且共享内存数组的大小为空（但是不可声明为指针）。

下面为使用共享内存进行归约求和程序的完整代码，包括了以上两种核函数：

typedef float real;
 
const int NUM_REPEATS = 100;
const int N = 100000000;
const int M = sizeof(real) * N;
const int BLOCK_SIZE = 128;
 
void timing(real *h_x, real *d_x, const int method);
 
int main(void) {
    real *h_x = (real *)malloc(M);
    for (int n = 0; n < N; ++n) {
        h_x[n] = 1.01;
    }
    real *d_x;
    CHECK(musaMalloc(&d_x, M));
 
    printf("\nUsing static shared memory:\n");
    timing(h_x, d_x, 0);
    printf("\nUsing dynamic shared memory:\n");
    timing(h_x, d_x, 1);
 
    free(h_x);
    CHECK(musaFree(d_x));
    return 0;
}
 
void __global__ reduce_shared(real *d_x, real *d_y) {
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int n = bid * blockDim.x + tid;
    __shared__ real s_y[128];
    s_y[tid] = (n < N) ? d_x[n] : 0.0;
    __syncthreads();
 
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            s_y[tid] += s_y[tid + offset];
        }
        __syncthreads();
    }
 
    if (tid == 0) {
        d_y[bid] = s_y[0];
    }
}
 
void __global__ reduce_dynamic(real *d_x, real *d_y) {
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int n = bid * blockDim.x + tid;
    extern __shared__ real s_y[];
    s_y[tid] = (n < N) ? d_x[n] : 0.0;
    __syncthreads();
 
    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            s_y[tid] += s_y[tid + offset];
        }
        __syncthreads();
    }
 
    if (tid == 0) {
        d_y[bid] = s_y[0];
    }
}
 
real reduce(real *d_x, const int method) {
    int grid_size = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    const int ymem = sizeof(real) * grid_size;
    const int smem = sizeof(real) * BLOCK_SIZE;
    real *d_y;
    CHECK(musaMalloc(&d_y, ymem));
    real *h_y = (real *)malloc(ymem);
 
    switch (method) {
        case 0:
            reduce_shared<<<grid_size, BLOCK_SIZE>>>(d_x, d_y);
            break;
        case 1:
            reduce_dynamic<<<grid_size, BLOCK_SIZE, smem>>>(d_x, d_y);
            break;
        default:
            printf("Error: wrong method\n");
            exit(1);
            break;
    }
 
    CHECK(musaMemcpy(h_y, d_y, ymem, musaMemcpyDeviceToHost));
 
    real result = 0.0;
    for (int n = 0; n < grid_size; ++n) {
        result += h_y[n];
    }
 
    free(h_y);
    CHECK(musaFree(d_y));
    return result;
}
 
void timing(real *h_x, real *d_x, const int method) {
    real sum = 0;
 
    float total_time;
    for (int repeat = 0; repeat < NUM_REPEATS; ++repeat) {
        CHECK(musaMemcpy(d_x, h_x, M, musaMemcpyHostToDevice));
 
        musaEvent_t start, stop;
        CHECK(musaEventCreate(&start));
        CHECK(musaEventCreate(&stop));
        CHECK(musaEventRecord(start));
        musaEventQuery(start);
 
        sum = reduce(d_x, method);
 
        CHECK(musaEventRecord(stop));
        CHECK(musaEventSynchronize(stop));
        float elapsed_time;
        CHECK(musaEventElapsedTime(&elapsed_time, start, stop));
 
        if (repeat > 0) {
            total_time += elapsed_time;
        }
 
        CHECK(musaEventDestroy(start));
        CHECK(musaEventDestroy(stop));
    }
    printf("average time = %f.\n", total_time / NUM_REPEATS);
 
    printf("sum = %f.\n", sum);
}

9.5.2. 优化实例2 - GEMM

矩阵乘法的例子

9.5.2.1. 概要

> 以下举例均是行主序，尺寸(行，列)

计算二维矩阵A和B的乘积C

• Matrix A : shape (m, k)
• Matrix B : shape (k, n)
• Matrix C : shape (m, n)

/*
    row_major, shape(row, col)
    output(m, n)
    input_a(m, k)
    input_b(k, n)
*/
__global__ void mat_mul_naive(float *output, const float *input_a, const float *input_b, int m, int n, int k)
{
    int tx = blockIdx.x * blockDim.x + threadIdx.x;
    int ty = blockIdx.y * blockDim.y + threadIdx.y;
    if (ty < m && tx < n)
    {
        float c = 0;
        for (int i = 0; i < k; ++i)
        {
            c += input_a[ty * k + i] * input_b[i * n + tx];
        }
        output[ty * n + tx] = c;
    }
}

Naive的实现过于简单，仅满足功能性实现，不符合追求高性能的需求，了解到其访存开销过大，计算占比过小后我们可以改良优化。

9.5.2.2. 基本思路

1. 将矩阵C换分为若干个块的子矩阵，由每个线程块计算一个子矩阵Block_C
   矩阵C的m和n维度分别按照tile_m和tile_n大小进行划分，即子矩阵块
   Block_C : shape (tile_m, tile_n)
   则矩阵A和矩阵B也自然被划分为：
   Block_A : shape (tile_m, k)
   Block_B : shape (k, tile_n)

2. 将线程划分为若干线程块，每块内每条线程计算Block_C中的N个元素，整体完成子矩阵的计算
   线程块总共有thread_per_block个线程， 每个线程处理N个元素， 最终完成tile_m*tile_n个元素
   因为共享内存大小有限， 多数情况下没法将Block_A和Block_B一次性加载完全， 所以关于维度k也可以进行进一步的的划分tile_k：
   Block_A : shape (tile_m, tile_k)
   Block_B : shape (tile_k, tile_n)

其中关键点在于映射好矩阵块和线程块的关系，包括内存的划分和搬运，以及数据的计算和写回。从全局内存加载对应的矩阵块到共享内存，再到寄存器，然后每个线程计算相应划分的结果，最终写回全局内存，完成矩阵块的乘积，进而得到整个矩阵乘积的结果。

9.5.2.3. 具体实现

如图所示，A矩阵和C矩阵行数相同，B矩阵和C矩阵的列相同，各维度对应符合矩阵乘法的要求，根据tile_m, tile_n和tile_k对矩阵进行分块，得到Block_A和Block_B，二者在K维度累计计算后得到矩阵块的乘积结果Block_C。

以子矩阵是正方块，每个线程处理一个元素为例(thread256_tile16x16)：
n = 1
tile_m = tile_n = tile_k = 16
则每个线程块包含256个线程，每个块的线程数量是slot大小的倍数而并且低于线程块的最大线程数限制。
每次加载对应的block的数据到共享内存，通过内部循环累加得到一个元素的结果。如下
设备代码

/*
    device code
    row_major, shape(row, col)
    output(m, n)
    input_a(m, k)
    input_b(k, n)
*/
__global__ void mat_mul_tile(float *output, const float *input_a, const float *input_b, int m, int n, int k)
{
    constexpr int TILE_M = 16;
    constexpr int TILE_N = 16;
    constexpr int TILE_K = 16;

    // shared memory for the block of input
    __shared__ float smem_a[TILE_M][TILE_K];
    __shared__ float smem_b[TILE_K][TILE_N];

    // block index
    int bx = blockIdx.x;
    int by = blockIdx.y;

    // thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    // the base address for the block of input
    const float *__restrict ptr_a_base = (const float *)(input_a + by * TILE_M * k);
    const float *__restrict ptr_b_base = (const float *)(input_b + bx * TILE_N);

    // the index of output
    int row = by * TILE_M + ty;
    int col = bx * TILE_N + tx;

    // the guard of input
    bool valid_input_a = row < m && tx < k;
    bool valid_input_b = ty < k && col < n;

    // the hot loop
    float reg_c = 0.0f;
    for (int i = 0; i < k + TILE_K - 1; i += TILE_K)
    {
        // load the tile of matrices from global memory to shared memory
        smem_a[ty][tx] = valid_input_a && i + tx < k ? ptr_a_base[k * ty + i + tx] : 0;
        smem_b[ty][tx] = valid_input_b && i + ty < k ? ptr_b_base[n * (i + ty) + tx] : 0;
        // synchronize to make sure that the loading is done
        __syncthreads();

        // multiply the two tiles of matrices
        for (int j = 0; j < TILE_K; ++j)
        {
            // load the element from shared memory to register
            float reg_a = smem_a[ty][j];
            float reg_b = smem_b[j][tx];
            // compute
            reg_c += reg_a * reg_b;
        }
        // synchronize to make sure that the computation is done before the next iteration
        __syncthreads();
    }

    // store the result from register to global memory
    output[row * n + col] = reg_c;
}

主机代码

// API_CHECK(expr) is the macro of musa api check function
void mat_mul(float *h_output, float *h_input_a, float *h_input_b, int m, int n, int k)
{
    // allocate matrix_a
    float *d_input_a;
    API_CHECK(musaMalloc(&d_input_a, sizeof(float) * m * k));

    // allocate matrix_b
    float *d_input_b;
    API_CHECK(musaMalloc(&d_input_b, sizeof(float) * k * n));

    // allocate matrix_c
    float *d_output;
    API_CHECK(musaMalloc(&d_output, sizeof(float) * m * n));

    // transfer input data to device memory
    API_CHECK(musaMemcpy(d_input_a, h_input_a, sizeof(float) * m * k, musaMemcpyHostToDevice));
    API_CHECK(musaMemcpy(d_input_b, h_input_b, sizeof(float) * k * n, musaMemcpyHostToDevice));

    // compute the execution configuration
    int tile_m = 16;
    int tile_n = 16;
    int grid_x = (n + tile_n - 1) / tile_n;
    int grid_y = (m + tile_m - 1) / tile_m;
    // launch the device kernel
    mat_mul_tile<<<dim3(grid_x, grid_y, 1), dim3(tile_n, tile_m, 1)>>>(d_output, d_input_a, d_input_b, m, n, k);
    
    // transfer the result from device to host
    API_CHECK(musaMemcpy(h_output, d_output, sizeof(float) * m * n, musaMemcpyDeviceToHost));
    
    // free memory
    API_CHECK(musaFree(d_input_a));
    API_CHECK(musaFree(d_input_b));
    API_CHECK(musaFree(d_output));
}

代码解释
代码包含两部分和两个函数：

• 主机代码 mat_mul() 负责在主机端调用设备kernel
• 设备代码 mat_mul_tile() 负责设备上矩阵乘法的kernel

主要以设备代码解释:
(1). 首先根据划分的块申请共享内存
(2). 根据块和线程的分布确定索引坐标等
(3). 加载矩阵块数据到共享内存, 使用同步确认线程块中所有线程加载完毕
(4). 从共享内存加载数据到寄存器, 线程块中每个线程计算一个结果, 使用同步确保所有线程完成计算
(5). 在K维度循环(3)， (4)步骤累加
(6). 最终将结果写回到全局内存完成

因为快速的共享内存的使用， 矩阵A和B从全局内存读取 (k / tile_k) 次， 降低了访存的开销。 但是这只能将访存代价从几百cycle降低到几十cycle，并没有改变问题的本质。问题的关键在于主体循环由两条 Load 指令与一条 FMA 指令构成，计算指令只占总体的 1/3，计算访存比过低，最终导致了访存延迟不能被隐藏，从而性能不理想。 所以我们可以提升计算访存比，每个线程处理更多的元素，进一步提升性能，从而接近理论峰值。

以128x128分块为例，每个线程处理64个元素，为了方便，假设m，n，k均是4的倍数(thread256_tile128x128):
n = 64
tile_k = 4
tile_m = tile_n = 128

每个线程块依然包含256个线程，为了能够合并访问共享内存，还是将线程划为二维16x16，同时使用向量指令加载，要将矩阵A转置存储，利用外积计算矩阵乘，减少访存指令数量。因为利用向量指令，计算访存比可以达到64/4，足够隐藏访存延迟。实际情况因为共享内存大小和寄存器的使用量限制，不同的线程块大小和每条线程计算的结果个数也不太同，需要根据硬件信息合理划分。在此我们以每个线程计算64个元素为例（本例没有添加double buffer，可进一步利用ping-pong交换节省同步开销）

__global__ void mat_mul_pseudo(float *output, const float *input_a, const float *input_b, int m, int n, int k)
{
    constexpr int tile_size = 128;
    constexpr int TILE_K = 4;

    // ping pong switch is optional
    // shared memory
    __shared__ float __attribute__((aligned(16))) smem_a[TILE_K * tile_size];
    __shared__ float __attribute__((aligned(16))) smem_b[TILE_K * tile_size];

    // register
    float4 load_reg_a[2];
    float4 load_reg_b[2];
    float4 reg_a[2][2];
    float4 reg_b[2][2];
    float4 reg_c[8][2] = {{float4(0.f, 0.f, 0.f, 0.f)}};

    // offset index compute
    int load_offset_a = ...;
    int load_offset_b = ...;
    int store_offset_a = ...;
    int store_offset_b = ...;

    // load first block from global memory to shared memory
    load_gmem_to_reg(input_a, load_offset_a, load_reg_a);
    load_gmem_to_reg(input_b, load_offset_b, load_reg_b);
    store_reg_to_smem_with_transpose(load_reg_a, store_offset_a, smem_a);
    store_reg_to_smem(load_reg_b, store_offset_b, smem_b);
    __syncthreads();

    // load first register from shared memory
    load_smem_to_reg(smem_a, 0, reg_a[0]);
    load_smem_to_reg(smem_b, 0, reg_b[0]);

    for (int i = TILE_K; i < k + TILE_K; i += TILE_K)
    {
        // load data from global memory
        load_gmem_to_reg(input_a, load_offset_a + i, load_reg_a);
        load_gmem_to_reg(input_b, load_offset_b + i, load_reg_b);

        for (int j = 0; j < TILE_K - 1; j++)
        {
            // load data from shared memory to register for next iteration
            load_smem_to_reg(smem_a, j + 1, reg_a[(j + 1) % 2]);
            load_smem_to_reg(smem_b, j + 1, reg_b[(j + 1) % 2]);
            // compute matrix multiply accumulate 8x8
            mma8x8(reg_c, reg_a[j % 2], reg_c[j % 2]);
        }

        if (i < k)
        {
            // make sure that the computation is done before the next iteration
            __syncthreads();
            // store data to shared memory before the next iteration
            store_reg_to_smem_with_transpose(load_reg_a, store_offset_a, smem_a);
            store_reg_to_smem(load_reg_b, store_offset_b, smem_b);
            __syncthreads();
        }

        // load data from shared memory to register before the next iteration
        load_smem_tile_to_reg(smem_a, 0, a_reg[0]);
        load_smem_tile_to_reg(smem_b, 0, b_reg[0]);
        // compute the last matrix multiply accumulate 8x8
        mma8x8(reg_c, a_reg[1], b_reg[1]);
    }
    // store the result register to global memory
    store_reg_to_gmem(output, reg_c);
}

由此我们获得了一个接近充分优化的矩阵乘法kernel，在MT GPU上实测，相较之前的之前的实现，更高的计算访存比和更合理的线程束划分使得性能得到质的提升。

9.5.3. 优化实例3 - CV

> 敬请期待后续版本更新