# CUDA 矩阵乘法终极优化指南

2021年09月15日 阅读数：3

## CUDA 矩阵乘法优化手段详解

### Naive 实现的分析：到底差在哪里？

``````__global__ void matrixMul(const float *A, const float *B, float *C,
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 += A[ty * K + i] * B[i * N + tx];
}
C[ty * N + tx] = c;
}
}
``````

``````    float c[4][4] = {{0}};
float a_reg[4];
float b_reg[4];
for(int i = 0; i < K; i += TILE_K){
// transfer tile from global mem to shared mem
#pragma unroll
for(int j = 0; j < TILE_K; ++j) {
// load tile from shared mem to register
// compute matrix multiply accumulate 4x4
mma4x4(a_reg, b_reg, c)；
}
}
``````

### 极致的访存优化

``````    int tx = threadIdx.x % 16;
int ty = threadIdx.x / 16;
``````

``````    #define TILE_K 16
__shared__ float4 smemA[2][TILE_K * 128 / 4];
__shared__ float4 smemB[2][TILE_K * 128 / 4];
float4 c[8][2] = {{make_float4(0.f, 0.f, 0.f, 0.f)}};
float4 ldg_a_reg[2];
float4 ldg_b_reg[2];
float4 a_reg[2][2];
float4 b_reg[2][2];

// transfer first tile from global mem to shared mem

store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[0]);
store_reg_to_smem_tile(ldg_b_reg, 0, smemB[0]);

// load first tile from shared mem to register

int write_stage_idx = 1; //ping pong switch
do {
i += TILE_K;
// load next tile from global mem

int load_stage_idx = write_stage_idx ^ 1;

#pragma unroll
for(int j = 0; j < TILE_K - 1; ++j) {
// load next tile from shared mem to register
// compute matrix multiply accumulate 8x8
mma8x8(a_reg[j % 2], b_reg[j % 2], c)；
}

if(i < K) {
// store next tile to shared mem
store_reg_to_smem_tile_transpose(ldg_a_reg, 0, smemA[write_stage_idx]);
store_reg_to_smem_tile(ldg_b_reg, 0, smemB[write_stage_idx]);
// use double buffer, only need one sync
// switch
write_stage_idx ^= 1;
}

// load first tile from shared mem to register of next iter
// compute last tile mma 8x8
mma8x8(a_reg[1], b_reg[1], c)；
} while (i < K);

store_c(c, C);
``````

## 超越 cublas：使用 SASS 调优 Kernel

### 寄存器 Bank conflict

maxas 中的 Maxwell SGEMM SASS Kernel 为了缓解寄存器 Bank conflict，就对参与 FFMA 计算的寄存器作了精巧的分配（参考 maxas 的 SGEMM 文档），以下图所示：

### Register Reuse

``````# Maxwell GPU
FFMA R2, R64.reuse, R73, R2; # R64 进入 Reuse Cache
FFMA R3, R64.reuse, R72, R3; # R64 从 Reuse Cache 中获取，避免与 R72 冲突
``````

## 总结

bot 说：文章都看完啦~ 是否有兴趣加入到深度学习框架开发中来？
Megengine 团队现正火热招聘中！期待你的加入~

【框架开发工程师（C++）】

1. 负责旷视核心深度框架 MegEngine 的设计，演进，实现，维护和优化
2. 优化 MegEngine 在各个计算平台（CUDA / Arm / x86 等）上的性能
3. 持续探索改进深度学习框架的先进优化方法（例如图优化，自动代码生成，超低 bit 量化，稀疏化等）

1. 1-3 年的性能优化经验（X86，CUDA，ARM，OpenCL 等）
2. 有良好的数据结构与算法功底，可以熟练使用 C++ 语言编写较复杂的算法
3. 对深度学习和深度学习框架（Caffe，Tensorflow，PyTorch 等）有基本了解者优先
4. 有复杂系统设计经验者优先