CUDA 编程（进阶篇）
GEMM 优化本身是一个非常值得讨论的课题，其优化也涉及 GPU 中优化的大多数常用的技巧。这部分以解析知乎大佬有了琦琦的棍子文章中的代码进行解读，也作为代码阅读笔记梳理整个思路
General Matrix Multiplication
首先，其优化技巧分块计算、shared memory 的多次利用、register 的多次利用以及各种 bank 的 conflict 解决，有的甚至会涉及到汇编层面的优化，这里有些技巧在基础篇已经讲过，就不再赘述了。
其次，简单叙述一下优化的思路，主要的思路就是对矩阵进行分块计算，不同 block 负责计算出 C 中的不同部分，同时在 block 内又让不同线程负责不同部分，这里面为了能多次利用 shared memory，需要进行多次循环，因此在 block 内有多次大循环，在大循环内又有每个线程中的多次小循环。因为涉及到把数据不断搬到 shared memory，所以作者设计了预取 prefetch 的做法，这样做可以掩盖 io 的 latency，因此也要设计哪些线程搬运哪些数据。由于可能在访问 shared memory 的时候有 bank conflict，所以也要设计哪些线程访问哪些内存
分块计算的思路
首先如下图，对 C 进行分块：
由图可知，C 被分为 MxN/BlOCK_SIZE_M/BlOCK_SIZE_N 块，每块的大小为高 BLOCK_SIZE_M，宽 BLOCK_SIZE_N，每一块对应 A 中的相应行，对应 B 中相应列。其中一个块交给一个 block 来计算，在本次示例中，A、B、C 都是 2048 的方阵，每块的大小为 128 的方阵，因此需要 256 个 block。
对于每个 block 而言，实际上是 A 中的某些行形成的矩阵（Ai）和 B 中的某些列形成的矩阵（Bi）进行矩阵乘法得到的。如上图所示，我们可以对 Ai 中的按列进行再一次划分（Aij）和 Bi 中按照相同的规则按行再次进行划分（Bij），其实可以发现 Ai 中的第 i 列只会和 Bi 中的第 i 行进行相应的运算，其实运算结果的矩阵就是 128 的方阵，所以其实 Ai 和 Bi 的矩阵乘法结果也可以看成多个 128 方阵的累加结果。对于一个 block 来说要进行 K/BLOCK_SIZE_K=8 次循环，然后将结果进行点对点叠加得到
对于每个 Aij 和 Bij 进行矩阵乘法的时候也不是直接按传统的 Aij 的某一行与 Bij 的某一列进行向量相乘得到最终结果的一个元素来算的，我们每次只计算 Aij 中的第 m 列和 Bi 中的第 m 行的运算结果。这实际上是有考虑 shared memory 的 bank conflict，并且虽然我们设定每个线程在每次小迭代中只计算 Aij 中的 8 个元素和 Bij 中的 8 个元素运算的结果，然而这 8 个元素并不是紧挨着的，如上如可知，我们开了 256 个线程来计算这 128x128 的矩阵，每个线程负责 64 个元素。256 个线程一共 8 个 wrap，我们需要使得每个 wrap 内不发生 bank conflict。上图展示了这个 8 个 wrap 负责的位置，以及 wrap 内线程的负责情况。下面先详细介绍一下 bank conflict：
Bank Conflict
在 shared memory 中的连续的内存被放入 32 个 bank 中，通常每 4 Bytes 一个 bank。如上图所示，一个 float 或者一个 int 占某一个 bank 中的一格，因此可以将一个 bank 理解为一个内存块，一个 bank 是一列，32 个 bank 组成 shared memory。一般 shared memory 的大小是 64 KB，也就是说一个 bank 最多有 2000 Bytes（500 行）。当我们给一个 wrap 中线程分配访问 shared memory 的地址时，也就是分配哪一个线程访问哪一个 bank 中的哪一格。什么是 bank conflict？当 同一 wrap 中 的 不同线程，访问 同一 bank 中的 不同地址 ，此时不能并行访问，只能串行，所以会导致速度下降。shared memory 中的访问可能性：
- 不同线程，访问不同 bank；这种情况是最希望的，因为可以并行访问
- 不同线程，访问同一 bank 中的同一地址（同一格），这种情况 gpu 会自动广播给所有线程，因此不影响速度
- 不同线程，访问同一 bank 中的不同地址（比如有 64 个线程，一个线程访问 4 Bytes，那么第 0 线程访问的 0，第 32 线程访问的 32，他们都属于 bank 0）这时候就会 conflict
为何不讨论不同 wrap 之间的 conflict 呢？因为不同 wrap 不能同时访问 SM 中的 shared memory。关于如何解决呢？ 目前一般有两个办法：
- 对 shared memory 开辟的长度进行改变，比如 32x6，对 0 线程而言 ，tid x 6 / 32，那么 16 线程就会发现冲突。如果将其改为 32 x 7 就不会了
- 对线程访问的地址进行设计，使其不会冲突
最后我们基于他的代码进行分析和学习（源代码注释也挺清楚的了）：
template <
const int BLOCK_SIZE_M, // C 中每个分块的高
const int BLOCK_SIZE_K, // A 中每次需要搬到共享内存块的宽
const int BLOCK_SIZE_N, // C 中每个分块的宽
const int THREAD_SIZE_Y, // C 中每个分块中一个线程需要计算的块的宽
const int THREAD_SIZE_X, // C 中每个分块中一个线程需要计算的块的高
const bool ENABLE_DOUBLE_BUFFER // 是否开启预取
>
__global__ void Sgemm(
float * __restrict__ A,
float * __restrict__ B,
float * __restrict__ C,
const int M,
const int N,
const int K) {
// Block 索引
int bx = blockIdx.x;
int by = blockIdx.y;
// Thread 索引
int tx = threadIdx.x;
int ty = threadIdx.y;
// 一个 block 线程的总数
const int THREAD_X_PER_BLOCK = BLOCK_SIZE_N / THREAD_SIZE_X;
const int THREAD_Y_PER_BLOCK = BLOCK_SIZE_M / THREAD_SIZE_Y;
const int THREAD_NUM_PER_BLOCK = THREAD_X_PER_BLOCK * THREAD_Y_PER_BLOCK;
// 线程在当前 block 中的索引
const int tid = ty * THREAD_X_PER_BLOCK + tx;
// 为搬运 A、B 矩阵，开了两份共享内存
__shared__ float As[2][BLOCK_SIZE_K][BLOCK_SIZE_M];
__shared__ float Bs[2][BLOCK_SIZE_K][BLOCK_SIZE_N];
// 在寄存器上开辟了记录累积结果以便把最终结果搬到 C 中
float accum[THREAD_SIZE_Y][THREAD_SIZE_X];
#pragma unroll
for(int i=0; i<THREAD_SIZE_Y; i++){
#pragma unroll
for(int j=0; j<THREAD_SIZE_X; j++){
accum[i][j]=0.0;
}
}
// 当前线程为搬运共享内存到寄存器上所开辟的内存
float frag_a[2][THREAD_SIZE_Y];
float frag_b[2][THREAD_SIZE_X];
// 这里计算每个线程每次搬运 4 个 float 需要搬运的次数（从 global 搬到 shared 上）
const int ldg_num_a = BLOCK_SIZE_M * BLOCK_SIZE_K / (THREAD_NUM_PER_BLOCK * 4);
const int ldg_num_b = BLOCK_SIZE_K * BLOCK_SIZE_N / (THREAD_NUM_PER_BLOCK * 4);
float ldg_a_reg[4*ldg_num_a];
float ldg_b_reg[4*ldg_num_b];
// 搬运一行需要的线程数
const int A_TILE_THREAD_PER_ROW = BLOCK_SIZE_K / 4;
const int B_TILE_THREAD_PER_ROW = BLOCK_SIZE_N / 4;
// 计算当前线程应该从哪一行哪一列开始搬运
const int A_TILE_ROW_START = tid / A_TILE_THREAD_PER_ROW;
const int B_TILE_ROW_START = tid / B_TILE_THREAD_PER_ROW;
const int A_TILE_COL = tid % A_TILE_THREAD_PER_ROW * 4;
const int B_TILE_COL = tid % B_TILE_THREAD_PER_ROW * 4;
// 总体搬运一次需要跨过的行数 -- stride
const int A_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / A_TILE_THREAD_PER_ROW;
const int B_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / B_TILE_THREAD_PER_ROW;
// 指向属于该 block 那个块的首地址
A = &A[(BLOCK_SIZE_M * by)* K];
B = &B[BLOCK_SIZE_N * bx];
// 需要该线程搬运共享内存的索引（从 shared 搬到 register 上）
const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int a_tile_index = warp_id/2*16 + lane_id/8*4;
const int b_tile_index = warp_id%2*32 + lane_id%8*4;
上面这份代码主要对需要的 index 进行了计算，其中
// 需要该线程搬运共享内存的索引（从 shared 搬到 register 上）
const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int a_tile_index = warp_id/2*16 + lane_id/8*4;
const int b_tile_index = warp_id%2*32 + lane_id%8*4;
这段代码的作用就是为了避免 bank conflict ，其效果如讲解分块时的图是一致的。首先，进行了计算该线程属于哪一个 wrap，在 wrap 中的局部 id 是多少，然后计算去共享内存搬运数据时从哪一个地址开始，举例说明：对于 A 矩阵需要搬运的共享内存而言，根据 wrap id 来确定大的位置（即这一片都是这个 wrap 需要搬运的数据，可参考前面分块画那个示意图）warp_id/2*16 得知每行两个 wrap，排完两个 wrap，跳（stride）16 个 foat。从 lane_id/8*4 可以看出 在 wrap 内 每行 8 个线程，也就是一共 4 行，4 行跳了 16 float，所以一个线程搬运 A 中的 4 个 float 刚好可以用 float4 这样可以访存对齐。由于开 256 个线程，一共 8 个 wrap 每行两个，一共 4 行，那才 64 float。但 A （转置后）一行有 128 个 float，所以需要搬运两次，因此一个线程搬运 8 个 float。
需要的 index 计算完成后，接下来就是数据的搬运了
// 首先将第一次大迭代相应的数据 global memory 搬运至 shared memory
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_M ; i += A_TILE_ROW_STRIDE) {
int ldg_index = i / A_TILE_ROW_STRIDE * 4;
FETCH_FLOAT4(ldg_a_reg[ldg_index]) = FETCH_FLOAT4(A[OFFSET(
A_TILE_ROW_START + i, // row
A_TILE_COL, // col
K)]);
As[0][A_TILE_COL][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index];
As[0][A_TILE_COL+1][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+1];
As[0][A_TILE_COL+2][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+2];
As[0][A_TILE_COL+3][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+3];
}
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_K; i += B_TILE_ROW_STRIDE) {
FETCH_FLOAT4(Bs[0][B_TILE_ROW_START + i][B_TILE_COL]) = FETCH_FLOAT4(B[OFFSET(
B_TILE_ROW_START + i, // row
B_TILE_COL, // col
N )]);
}
__syncthreads();
// 再把 shared memory 中的第一次小迭代数据搬到 register 中
FETCH_FLOAT4(frag_a[0][0]) = FETCH_FLOAT4(As[0][0][a_tile_index]);
FETCH_FLOAT4(frag_a[0][4]) = FETCH_FLOAT4(As[0][0][a_tile_index + 64]);
FETCH_FLOAT4(frag_b[0][0]) = FETCH_FLOAT4(Bs[0][0][b_tile_index]);
FETCH_FLOAT4(frag_b[0][4]) = FETCH_FLOAT4(Bs[0][0][b_tile_index + 64]);
可以注意到上一段代码，最后搬到 register 确实搬了两次。首先，这段代码将第一次大迭代需要的数据存入了 shared memory 里面，然后把第一次小迭代需要的数据存入了 register，其中：
FETCH_FLOAT4：是利用 cuda 中的类型 float4 而设置的宏，很明显一次可以取 4 个 float，可以避免访存不对齐的问题；ldg_a_reg：显然要想从 A 中搬运数据到存 A 数据的 shared memory 需要转置，此时不能直接使用FETCH_FLOAT4
接下来就是正式计算的代码了
// 作为预取数据的数组 index
int write_stage_idx = 1;
int tile_idx = 0;
do{
// 大迭代中每次的 stride 为 BLOCK_SIZE_K
tile_idx += BLOCK_SIZE_K;
// 先将下一次大迭代的 global memory 中的数据考到 register 中
if(tile_idx< K){
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_M ; i += A_TILE_ROW_STRIDE) {
int ldg_index = i / A_TILE_ROW_STRIDE * 4;
FETCH_FLOAT4(ldg_a_reg[ldg_index]) = FETCH_FLOAT4(A[OFFSET(
A_TILE_ROW_START + i, // row
A_TILE_COL + tile_idx, // col
K )]);
}
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_K; i += B_TILE_ROW_STRIDE) {
int ldg_index = i / B_TILE_ROW_STRIDE * 4;
FETCH_FLOAT4(ldg_b_reg[ldg_index]) = FETCH_FLOAT4(B[OFFSET(
tile_idx + B_TILE_ROW_START + i, // row
B_TILE_COL, // col
N )]);
}
}
// 作为加载数据的数组的 index write 为 0 load 就为 1
int load_stage_idx = write_stage_idx ^ 1;
#pragma unroll
for(int j=0; j<BLOCK_SIZE_K - 1; ++j){
// 第一次大迭代中 load_stage_idx 就为 0
// 因为在上部分代码中，我们已经进行了数据的预取
// 这里把下一次小迭代的需要数据放入 register 中
FETCH_FLOAT4(frag_a[(j+1)%2][0]) = FETCH_FLOAT4(As[load_stage_idx][(j+1)][a_tile_index]);
FETCH_FLOAT4(frag_a[(j+1)%2][4]) = FETCH_FLOAT4(As[load_stage_idx][(j+1)][a_tile_index + 64]);
FETCH_FLOAT4(frag_b[(j+1)%2][0]) = FETCH_FLOAT4(Bs[load_stage_idx][(j+1)][b_tile_index]);
FETCH_FLOAT4(frag_b[(j+1)%2][4]) = FETCH_FLOAT4(Bs[load_stage_idx][(j+1)][b_tile_index + 64]);
// 这里把第一次小迭代的数据拿出来进行计算
#pragma unroll
for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
accum[thread_y][thread_x] += frag_a[j%2][thread_y] * frag_b[j%2][thread_x];
}
}
}
if(tile_idx < K){
// 将刚才的预取到 register 中的数据，存到 shared momery 中
// 供下一次小迭代反复使用
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_M ; i += A_TILE_ROW_STRIDE) {
int ldg_index = i / A_TILE_ROW_STRIDE * 4;
As[write_stage_idx][A_TILE_COL][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index];
As[write_stage_idx][A_TILE_COL+1][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+1];
As[write_stage_idx][A_TILE_COL+2][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+2];
As[write_stage_idx][A_TILE_COL+3][A_TILE_ROW_START + i]=ldg_a_reg[ldg_index+3];
}
#pragma unroll
for ( int i = 0 ; i < BLOCK_SIZE_K; i += B_TILE_ROW_STRIDE) {
int ldg_index = i / B_TILE_ROW_STRIDE * 4;
FETCH_FLOAT4(Bs[write_stage_idx][B_TILE_ROW_START + i][B_TILE_COL]) = FETCH_FLOAT4(ldg_b_reg[ldg_index]);
}
__syncthreads();
// 下一次的预取索引就应该是上一次的加载索引
// 同理上面使用的预取索引就是下一次的加载索引 预取就是为了未来的加载嘛
write_stage_idx ^= 1;
}
// 把第一次小迭代的数据预取至寄存器
FETCH_FLOAT4(frag_a[0][0]) = FETCH_FLOAT4(As[load_stage_idx^1][0][a_tile_index]);
FETCH_FLOAT4(frag_a[0][4]) = FETCH_FLOAT4(As[load_stage_idx^1][0][a_tile_index + 64]);
FETCH_FLOAT4(frag_b[0][0]) = FETCH_FLOAT4(Bs[load_stage_idx^1][0][b_tile_index]);
FETCH_FLOAT4(frag_b[0][4]) = FETCH_FLOAT4(Bs[load_stage_idx^1][0][b_tile_index + 64]);
// 计算最后一次小迭代
#pragma unroll
for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
accum[thread_y][thread_x] += frag_a[1][thread_y] * frag_b[1][thread_x];
}
}
}while(tile_idx< K);
每次大迭代中最后一次小迭代放到搬完数据的最后来做，猜测这样在执行指令的时候就可以边取存下一次迭代的数据时，边执行计算。最后把每个线程寄存器上的结果搬回 C 中
const int c_block_row = a_tile_index;
const int c_block_col = b_tile_index;
//store C00 block
for(int i=0; i<4; i++){
FETCH_FLOAT4(C[OFFSET(
BLOCK_SIZE_M * by + c_block_row + i,
BLOCK_SIZE_N * bx + c_block_col,
N)]) = FETCH_FLOAT4(accum[i][0]);
}
//store C01 block
for(int i=0; i<4; i++){
FETCH_FLOAT4(C[OFFSET(
BLOCK_SIZE_M * by + c_block_row + i,
BLOCK_SIZE_N * bx + c_block_col + 64,
N)]) = FETCH_FLOAT4(accum[i][4]);
}
//store C10 block
for(int i=0; i<4; i++){
FETCH_FLOAT4(C[OFFSET(
BLOCK_SIZE_M * by + c_block_row + 64 + i,
BLOCK_SIZE_N * bx + c_block_col,
N)]) = FETCH_FLOAT4(accum[i+4][0]);
}
//store C11 block
for(int i=0; i<4; i++){
FETCH_FLOAT4(C[OFFSET(
BLOCK_SIZE_M * by + c_block_row + 64 + i,
BLOCK_SIZE_N * bx + c_block_col + 64,
N)]) = FETCH_FLOAT4(accum[i+4][4]);
}
Trick
- 首先分块：如何分给 block，block 内如何分给线程
- 然后搬运数据：大迭代的数据怎么搬运、小迭代的数据怎么搬运、利用 shared memory 和 register 提高搬运数据的效率；利用预取掩盖访问 latency
- 最后如何防止 bank conflict，对各个 wrap 应该访问的地址进行了设计
To Be Continued
🌸 February 9, 2023
Comments
Sign in with GitHub to comment, or click the bottom icon to chat with me.