gemm学习笔记

写在前面

经过了一个寒假，发现之前自己学习的gemm都快忘的差不多了，好记性还是不如烂笔头嘛，于是开个新专栏，记录下自己的学习

初始化

使用的环境是之前构建的docker，首先ssh连接上电脑（ssh siyuan@dell），然后在vscode中找到该docker并选择attach to terminal，然后输入 service start ssh,此后便可通过ssh siyuan@docker进入此docker

使用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
#pragma once
#include <stdio.h>

#define CHECK(call)                                   \
do                                                    \
{                                                     \
    const cudaError_t error_code = call;              \
    if (error_code != cudaSuccess)                    \
    {                                                 \
        printf("CUDA Error:\n");                      \
        printf("    File:       %s\n", __FILE__);     \
        printf("    Line:       %d\n", __LINE__);     \
        printf("    Error code: %d\n", error_code);   \
        printf("    Error text: %s\n",                \
            cudaGetErrorString(error_code));          \
        exit(1);                                      \
    }                                                 \
} while (0)
#include <stdio.h>

int main(int argc, char *argv[])
{
    int device_id = 0;
    if (argc > 1) device_id = atoi(argv[1]);
    CHECK(cudaSetDevice(device_id));

    cudaDeviceProp prop;
    CHECK(cudaGetDeviceProperties(&prop, device_id));

    printf("Device id:                                 %d\n",
        device_id);
    printf("Device name:                               %s\n",
        prop.name);
    printf("Compute capability:                        %d.%d\n",
        prop.major, prop.minor);
    printf("Amount of global memory:                   %g GB\n",
        prop.totalGlobalMem / (1024.0 * 1024 * 1024));
    printf("Amount of constant memory:                 %g KB\n",
        prop.totalConstMem  / 1024.0);
    printf("Maximum grid size:                         %d %d %d\n",
        prop.maxGridSize[0], 
        prop.maxGridSize[1], prop.maxGridSize[2]);
    printf("Maximum block size:                        %d %d %d\n",
        prop.maxThreadsDim[0], prop.maxThreadsDim[1], 
        prop.maxThreadsDim[2]);
    printf("Number of SMs:                             %d\n",
        prop.multiProcessorCount);
    printf("Maximum amount of shared memory per block: %g KB\n",
        prop.sharedMemPerBlock / 1024.0);
    printf("Maximum amount of shared memory per SM:    %g KB\n",
        prop.sharedMemPerMultiprocessor / 1024.0);
    printf("Maximum number of registers per block:     %d K\n",
        prop.regsPerBlock / 1024);
    printf("Maximum number of registers per SM:        %d K\n",
        prop.regsPerMultiprocessor / 1024);
    printf("Maximum number of threads per block:       %d\n",
        prop.maxThreadsPerBlock);
    printf("Maximum number of threads per SM:          %d\n",
        prop.maxThreadsPerMultiProcessor);

    return 0;
}
CPP

查看显卡的基本信息，得到信息如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
Device id:                                 0
Device name:                               NVIDIA GeForce GTX 1650
Compute capability:                        7.5
Amount of global memory:                   3.80554 GB
Amount of constant memory:                 64 KB
Maximum grid size:                         2147483647 65535 65535
Maximum block size:                        1024 1024 64
Number of SMs:                             16
Maximum amount of shared memory per block: 48 KB
Maximum amount of shared memory per SM:    64 KB
Maximum number of registers per block:     64 K
Maximum number of registers per SM:        64 K
Maximum number of threads per block:       1024
Maximum number of threads per SM:          1024
APACHE

gemm

我的一些问题

  1. 为什么使用列索引：

    个人理解是：

    cublas为了适配Fortran，默认使用了列优先存储 。在test.cu的实际调用cublasSgemm(err, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, dA, m, dB, k, &beta, dC_ref, m);中，根据api可推断确实使用的是列优先。
    当把 C / C++ 中行优先保存的二维数组 A 以一维形式输入 cublas 时，会被 cublas 理解为列优先存储。
    为了保证和cublas的一致性，且避免行列的混乱，统一使用了列主序，从kernel1.cu中的#define A(i,j) A[(i) + (j)*lda] #define B(i,j) B[(i) + (j)*ldb] #define C(i,j) C[(i) + (j)*ldc]中也可推断出kernel中的矩阵是列优先的。

  2. 通俗的计算逻辑：

    我认为可以分为大子块和小子块两个部分，每个grid完成大子块的计算，通常是一个32*32的矩阵，在大子块中，每个线程负责一小块矩阵的计算，这在kernel4以前都是1*1，在kernel5变成了4*1，kernel6变成了4*4。

源代码分析

原始实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
// naive version
__global__  __launch_bounds__(1024)
void mysgemm_v1(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x, ty = threadIdx.y;
    int bx = blockIdx.x, by = blockIdx.y;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    float tmp=0.;
    for (int k_count = 0; k_count<K; k_count++){
        tmp += A(tx, k_count) * B(k_count, ty);
    }
    C(tx,ty) = alpha * tmp + beta*C(tx,ty);
}

CPP

  • 宏定义

    lda是A矩阵的行数（M），ldb是b矩阵的行数（K）

    1
    2
    3
    #define A(i,j) A[(i) + (j)*lda]
    #define B(i,j) B[(i) + (j)*ldb]
    #define C(i,j) C[(i) + (j)*ldc]
    CPP

    由于是列索引，存储是A[0,0]- A[1,0]-A[2,0]这样先存储完一整列，再存储下一列，因此这样的define可以正确定位

  • 移位

    1
    2
    3
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    CPP

    每一个block负责32行A的子矩阵和32列B子矩阵的乘法，因此通过A = &A((bx<<5),0);
    B = &B(0,(by<<5));直接更改A和B的地址。对于C矩阵，C = &C((bx<<5),(by<<5));也是直接定位到此block计算对应的子矩阵位置

  • 计算

    1
    2
    3
    4
    5
    float tmp=0.;
    for (int k_count = 0; k_count<K; k_count++){
        tmp += A(tx, k_count) * B(k_count, ty);
    }
    C(tx,ty) = alpha * tmp + beta*C(tx,ty);
    CPP

    在每个block的内部，32*32个thread，每个线程各负责一个子矩阵中元素的计算。宏观上讲，该block中的第（tx,ty)个线程负责C（bx*32+tx,by*32+ty）元素的计算。

    实现全部计算。

引入共享内存

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa(i,j) sa[((i)<<5) + (j)]
#define sb(i,j) sb[((i)<<5) + (j)]
#define MS 32
#define NS 32
#define KS 32
// cache blocking version, without register-level data re-use
__global__  __launch_bounds__(1024)
void mysgemm_v2(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x, ty = threadIdx.y;
    int bx = blockIdx.x, by = blockIdx.y;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    __shared__ float sa[MS*KS];
    __shared__ float sb[KS*NS];
    float tmp=0.;
    for (int k_count = 0; k_count<K; k_count+=KS){
        sa(tx,ty)=A(tx,ty);
        sb(ty,tx)=B(tx,ty);
        A+=(lda<<5);B+=32;
        __syncthreads();
        for (int inner_k_count=0;inner_k_count<KS;inner_k_count++){
            tmp += sa(tx,inner_k_count) * sb(ty,inner_k_count);
        }
        __syncthreads();
    }
    C(tx,ty) = alpha * tmp + beta*C(tx,ty);
}
CPP

  • 宏定义

    1
    2
    #define sa(i,j) sa[((i)<<5) + (j)]
    #define sb(i,j) sb[((i)<<5) + (j)]
    CPP

    完成共享内存的宏定义，表示sa和sb使用行优先存储，与ABC都不同！

  • 共享内存

    1
    2
    __shared__ float sa[MS*KS];
    __shared__ float sb[KS*NS];
    CPP

    开了两块共享内存

  • 内存分配

    1
    2
    3
    for (int k_count = 0; k_count<K; k_count+=KS){
        sa(tx,ty)=A(tx,ty);
        sb(ty,tx)=B(tx,ty);}
    CPP

    此部分的分配主要为了方便后续计算，tmp += sa(tx,inner_k_count) * sb(ty,inner_k_count)能够按存储顺序读取sa和sb元素，减小读取的延迟。

后续的计算与kernel1保持一致。

省一个寄存器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa(i,j) sa[((i)<<5) + (j)]
#define sb(i,j) sb[((i)<<5) + (j)]
#define MS 32
#define NS 32
#define KS 32
// cache blocking version, without register-level data re-use
// save one living register ty.
__global__  __launch_bounds__(1024)
void mysgemm_v3(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row = tx&31, col = tx>>5;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    __shared__ float sa[MS*KS];
    __shared__ float sb[KS*NS];
    float tmp=0.;
    for (int k_count = 0; k_count<K; k_count+=KS){
        sa(row,col)=A(row,col);
        sb(col,row)=B(row,col);
        A+=(lda<<5);B+=32;
        __syncthreads();
        for (int inner_k_count=0;inner_k_count<KS;inner_k_count++){
            tmp += sa(row,inner_k_count) * sb(col,inner_k_count);
        }
        __syncthreads();
    }
    C(row,col) = alpha * tmp + beta*C(row,col);
}
CPP

与上一版本相比，区别仅仅在于使用int tx = threadIdx.x;int row = tx&31, col = tx>>5;说是省了一个寄存器，但是row和col变量的引入可能又占用了别的内存，具体效果有待用nsight验证。

另外sb4处的col、row指代的是B中的row和col，在后续sb4(col,inner_k_count)部分需要注意区分以免引起歧义

Memory coalescing

Memory coalescing on shared memory 是并行计算（如 GPU 编程）中的一种优化技术，主要目的是提高内存访问效率，尤其在处理共享内存时。在 GPU 等并行计算架构中，多个线程同时访问共享内存。如果这些访问能够合并为更少的内存事务，就能显著提升性能。

在kernel2和3中，就已经正确的使用了Memory coalescing，此部分的代码只是为了展示没有完成内存对齐的效果如何。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa4(i,j) sa4[((j)<<5) + (i)]
#define sb4(i,j) sb4[((j)<<5) + (i)]
#define MS 32
#define NS 32
#define KS 32
// cache blocking version, without register-level data re-use
// with memory coelascing on shared memory
__global__  __launch_bounds__(1024)
void mysgemm_v4(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row = tx&31, col = tx>>5;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    __shared__ float sa4[MS*KS];
    __shared__ float sb4[KS*NS];
    float tmp=0.;
    for (int k_count = 0; k_count<K; k_count+=KS){
        sa4(row,col)=A(row,col);
        sb4(col,row)=B(row,col);
        A+=(lda<<5);B+=32;
        __syncthreads();
        for (int inner_k_count=0;inner_k_count<KS;inner_k_count++){
            tmp += sa4(row,inner_k_count) * sb4(col,inner_k_count);
        }
        __syncthreads();
    }
    C(row,col) = alpha * tmp + beta*C(row,col);
}
CPP

使用micro kernel提升计算强度（Arithmetic Intensity）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa5(i,j) sa5[((j)<<5) + (i)]
#define sb5(i,j) sb5[((j)<<5) + (i)]
#define MS 32
#define NS 32
#define KS 32
// cache blocking version, without register-level data re-use
// with memory coelascing on shared memory
// more workloads per thread. 4x1 micro kernel.
__global__  __launch_bounds__(256)
void mysgemm_v5(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row1 = (tx&7)<<2, row2 = row1+1, row3 = row1+2, row4 = row1+3, col = tx>>3;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    __shared__ float sa5[MS*KS];
    __shared__ float sb5[KS*NS];
    float Cres[4] = {0., 0., 0., 0.};
    float b00;
    for (int k_count = 0; k_count<K; k_count+=KS){
        sa5(row1,col)=A(row1,col);
        sa5(row2,col)=A(row2,col);
        sa5(row3,col)=A(row3,col);
        sa5(row4,col)=A(row4,col);
        sb5(col,row1)=B(row1,col);
        sb5(col,row2)=B(row2,col);
        sb5(col,row3)=B(row3,col);
        sb5(col,row4)=B(row4,col);
        A+=(lda<<5);B+=32;
        __syncthreads();
        #pragma unroll
        for (int inner_k_count=0;inner_k_count<KS;inner_k_count++){
            b00 = sb5(col,inner_k_count);
            Cres[0] += sa5(row1,inner_k_count) * b00;
            Cres[1] += sa5(row2,inner_k_count) * b00;
            Cres[2] += sa5(row3,inner_k_count) * b00;
            Cres[3] += sa5(row4,inner_k_count) * b00;
        }
        __syncthreads();
    }
    C(row1,col) = alpha * Cres[0] + beta*C(row1,col);
    C(row2,col) = alpha * Cres[1] + beta*C(row2,col);
    C(row3,col) = alpha * Cres[2] + beta*C(row3,col);
    C(row4,col) = alpha * Cres[3] + beta*C(row4,col);
}
CPP

此kernel与前几版相比，区别在于之前的每个thread都计算C矩阵中1个元素的值，现在的thread需要计算4个，以此提高了计算强度。

具体到代码实现上，区别主要在于：

  1. row的确定

    1
    int row1 = (tx&7)<<2, row2 = row1+1, row3 = row1+2, row4 = row1+3, col = tx>>3;
    CPP

    此部分需要重点理解的是

    • tx&7：取 tx 的低 3 位（tx % 8），用于确定线程在 8x8 的子块中的行偏移。
    • (tx&7)<<2：将行偏移乘以 4，得到 row1，表示线程负责的起始行。
    • row2, row3, row4：线程负责的另外 3 行（row1 + 1, row1 + 2, row1 + 3）。
    • tx>>3：将 tx 右移 3 位（tx / 8），用于确定线程在 8x8 子块中的列偏移，赋值给 col。

    在完成此操作后，row1的值可能为0，4，8，12…28，row2为1，5…29，以此覆盖32行，对于列由于tx的范围为0…256，故col的范围为0…32。

    在此条件下，以tx=133为例，row1=（133%8）*4=20，col=（133/8）=16，因此tx=133的线程负责小矩阵块中（20，16）、(21,16)…(23,16)四个元素的计算

  2. 变量定义

    1
    float Cres[4] = {0., 0., 0., 0.};
    CPP

    原本的C变成了Cres，配合后面的

    1
    2
    3
    4
    5
    Cres[0] += sa5(row1,inner_k_count) * b00;
            Cres[1] += sa5(row2,inner_k_count) * b00;
            Cres[2] += sa5(row3,inner_k_count) * b00;
            Cres[3] += sa5(row4,inner_k_count) * b00;
        }
    CPP

    完成了单个线程的计算

  3. 共享内存赋值

    1
    2
    3
    4
    5
    6
    7
    8
    9
    sa5(row1,col)=A(row1,col);
        sa5(row2,col)=A(row2,col);
        sa5(row3,col)=A(row3,col);
        sa5(row4,col)=A(row4,col);
        sb5(col,row1)=B(row1,col);
        sb5(col,row2)=B(row2,col);
        sb5(col,row3)=B(row3,col);
        sb5(col,row4)=B(row4,col);

    CPP

    这块没什么特别的，只是把之前的赋值展开成了四个，后续计算同理

向量化存储/读取

好处：

  • 内存带宽利用率高：一次性加载或存储4个浮点数，减少了内存访问次数，提高了内存带宽的利用率。
  • 并行性增强：GPU的SIMD（单指令多数据）架构适合处理向量化数据，float4 可以更好地利用硬件的并行计算能力。
  • 减少指令开销：相比于逐个加载或存储浮点数，向量化操作可以减少指令的数量，从而降低指令开销。

float4定义：

1
2
3
4
5
6
struct float4 {
    float x; // 第一个浮点数
    float y; // 第二个浮点数
    float z; // 第三个浮点数
    float w; // 第四个浮点数
};
CPP
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa6(i,j) sa6[((j)<<5) + (i)]
#define sb6(i,j) sb6[((j)<<5) + (i)]
#define MS 32
#define NS 32
#define KS 32
// cache blocking version, without register-level data re-used
// with memory coelascing on shared memory
// more workloads per thread. 4x1 micro kernel.
// adopt vetorized load/store
__global__  __launch_bounds__(256)
void mysgemm_v6(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row1 = (tx&7)<<2, row2 = row1+1, row3 = row1+2, row4 = row1+3, col = tx>>3;
    A = &A((bx<<5),0);
    B = &B(0,(by<<5));
    C = &C((bx<<5),(by<<5));
    __shared__ float sa6[MS*KS];
    __shared__ float sb6[KS*NS];
    float4 Av, Bv, Cv, Cres;
    Cres.x = 0., Cres.y = 0., Cres.z = 0., Cres.w = 0.;
    float b00;
    for (int k_count = 0; k_count<K; k_count+=KS){
        Av = *((float4 *)(&A(row1,col)));
        Bv = *((float4 *)(&B(row1,col)));
        ((float4 *)sa6)[tx] = Av;
        sb6(col,row1)=Bv.x;
        sb6(col,row2)=Bv.y;
        sb6(col,row3)=Bv.z;
        sb6(col,row4)=Bv.w;
        A+=(lda<<5);B+=32;
        __syncthreads();
        #pragma unroll
        for (int inner_k_count=0;inner_k_count<KS;inner_k_count++){
            b00 = sb6(col, inner_k_count);
            Cres.x += sa6(row1,inner_k_count) * b00;
            Cres.y += sa6(row2,inner_k_count) * b00;
            Cres.z += sa6(row3,inner_k_count) * b00;
            Cres.w += sa6(row4,inner_k_count) * b00;
        }
        __syncthreads();
    }
    Cv = *((float4 *)(&C(row1,col)));
    Cres.x = alpha * Cres.x + beta * Cv.x;
    Cres.y = alpha * Cres.y + beta * Cv.y;
    Cres.z = alpha * Cres.z + beta * Cv.z;
    Cres.w = alpha * Cres.w + beta * Cv.w;
    *(float4 *)(&(C(row1,col))) = Cres;
}
CPP

与之前的核的区别在于使用了float4 Av, Bv, Cv, Cres;来向量化计算，核心代码：

1
2
3
4
5
6
7
Av = *((float4 *)(&A(row1,col)));
        Bv = *((float4 *)(&B(row1,col)));
        ((float4 *)sa6)[tx] = Av;
        sb6(col,row1)=Bv.x;
        sb6(col,row2)=Bv.y;
        sb6(col,row3)=Bv.z;
        sb6(col,row4)=Bv.w;
CPP

对于sa6，直接按照4个一组的顺序把A矩阵存入，sb则进行了转置，此处逻辑较为混乱，但整体思想很好理解，后续对Cres的计算更是非常简单。

增加线程工作量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa7(i,j) sa7[((j)<<6) + (i)]
#define sb7(i,j) sb7[((j)<<6) + (i)]
#define MS_7 64
#define NS_7 64
#define KS_7 16
//v1 += v2 * s3, vector scaling
#define vscal(v1, v2, s3)\
    v1.x+=v2.x*s3;\
    v1.y+=v2.y*s3;\
    v1.z+=v2.z*s3;\
    v1.w+=v2.w*s3;
//v1 = alpha * v2 + beta * v3, simd fma
#define simd_axpby(v1, alpha, v2, beta, v3)\
    v1.x=alpha*v2.x+beta*v3.x;\
    v1.y=alpha*v2.y+beta*v3.y;\
    v1.z=alpha*v2.z+beta*v3.z;\
    v1.w=alpha*v2.w+beta*v3.w;
#define vload(v1,addr)\
    v1 = *((float4 *)(addr));
#define vstore(addr,v1)\
    *((float4 *)(addr)) = v1;
// cache blocking version, without register-level data re-use
// with memory coelascing on shared memory
// more workloads per thread. 4x4 micro kernel.
// adopt vetorized load/store
__global__  __launch_bounds__(256)
void mysgemm_v7(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row_a = (tx&15)<<2, col_a = tx>>4;
    int row_b = (tx&3)<<2, col_b = tx>>2;
    int col_c = col_a<<2;
    int lda16 = lda<<4;
    A = &A((bx<<6),0);
    B = &B(0,(by<<6));
    C = &C((bx<<6),(by<<6));//the TB size is 64.
    __shared__ float sa7[1024];
    __shared__ float sb7[1024];
    float4 Av, Bv, Cv[4], Cres[4];
    memset(Cres, 0, sizeof(Cres));
    for (int k_count = 0; k_count<K; k_count+=KS_7){
        vload(Av, &A(row_a,col_a))
        vload(Bv, &B(row_b,col_b))
        ((float4 *)sa7)[tx] = Av;
        sb7(col_b,row_b)=Bv.x;
        sb7(col_b,row_b+1)=Bv.y;
        sb7(col_b,row_b+2)=Bv.z;
        sb7(col_b,row_b+3)=Bv.w;
        A+=lda16;B+=16;
        __syncthreads();
        #pragma unroll
        for (int inner_k_count=0;inner_k_count<KS_7;inner_k_count++){
            vload(Av, &sa7(row_a,inner_k_count))
            vload(Bv, &sb7(col_c,inner_k_count))
            vscal(Cres[0], Av, Bv.x)
            vscal(Cres[1], Av, Bv.y)
            vscal(Cres[2], Av, Bv.z)
            vscal(Cres[3], Av, Bv.w)
        }
        __syncthreads();
    }
    vload(Cv[0], &C(row_a,col_c))
    vload(Cv[1], &C(row_a,col_c+1))
    vload(Cv[2], &C(row_a,col_c+2))
    vload(Cv[3], &C(row_a,col_c+3))
    simd_axpby(Cres[0],alpha,Cres[0],beta,Cv[0])
    simd_axpby(Cres[1],alpha,Cres[1],beta,Cv[1])
    simd_axpby(Cres[2],alpha,Cres[2],beta,Cv[2])
    simd_axpby(Cres[3],alpha,Cres[3],beta,Cv[3])

    vstore(&C(row_a,col_c), Cres[0])
    vstore(&C(row_a,col_c+1), Cres[1])
    vstore(&C(row_a,col_c+2), Cres[2])
    vstore(&C(row_a,col_c+3), Cres[3])
}
CPP

DeepSeek版解析：

  1. 核心思想：增加每个线程的工作量

    背景：GPU 线程块（TB）的寄存器资源充足（假设 64K），而每个 TB 仅分配了 256 个线程。此时单个线程的寄存器使用量较低，存在资源浪费。
    优化策略：让每个线程计算更大的子矩阵（例如 4×4 的 C 子矩阵）。
    优势：
    寄存器级数据重用：每个线程在寄存器中保留更多中间计算结果，减少对共享内存或全局内存的访问次数。
    减少内存延迟：寄存器访问速度远快于共享内存和全局内存，从而降低计算延迟。
    提高计算强度（Arithmetic Intensity）：每个线程执行更多计算操作，减少内存访问与计算的比率。

  2. 参数调整：Ms、Ns 和 Ks

    原始参数：Ms = 32, Ns = 32, Ks = 32。
    新参数：Ms = 64, Ns = 64, Ks = 16。
    目标：
    保持共享内存（Shared Memory）的总消耗不变。
    通过增大 Ms 和 Ns，让每个线程块处理更大的数据块。
    通过减少 Ks，抵消 Ms 和 Ns 增大带来的共享内存压力。
    共享内存计算

    共享内存占用公式：Shared Memory Size = Ms × Ks + Ns × Ks。
    原始值：32×32 + 32×32 = 2048 单位（例如字节）。
    新值：64×16 + 64×16 = 2048 单位（保持不变）。
    意义：在共享内存资源受限的情况下，通过调整 Ks，允许每个线程块处理更大的输入矩阵块（Ms 和 Ns 增大）。

  3. 性能优势分析

    (1) 输入矩阵较大时的优势

    更多线程块（TBs）映射到流式多处理器（SMs）：
    增大 Ms 和 Ns 后，每个线程块处理的数据量增加，但 GPU 的 SMs 仍然可以调度足够的线程块来覆盖整个输入矩阵。
    例如：若输入矩阵尺寸为 4096×4096，增大 Ms 和 Ns 可以减少所需的线程块总数，同时保持高并行度。
    减少全局内存访问：
    每个线程块处理更多数据，减少了从全局内存加载数据的频率（数据局部性更好）。
    (2) 线程级优化

    更高的寄存器利用率：
    每个线程计算更大的子矩阵（如 4×4），充分利用寄存器资源，减少共享内存访问。
    减少线程同步开销：
    每个线程完成更多独立计算，减少线程间同步（如 __syncthreads()）的频率。

  4. 适用场景

    大规模矩阵乘法：当输入矩阵足够大时，参数调整的优化效果更明显。
    共享内存受限的 GPU 架构：通过平衡 Ms、Ns 和 Ks，避免共享内存成为瓶颈。

  5. 总结

    这段描述的核心是通过以下策略优化 GPU 矩阵乘法：

    增加单个线程的计算量（4×4 子矩阵），利用寄存器重用数据，减少内存访问。
    调整分块参数（Ms、Ns、Ks），在共享内存资源不变的前提下，扩大每个线程块处理的数据规模。
    提升计算强度，更好地利用 GPU 的并行计算能力，尤其适用于大规模矩阵运算。
    这种优化方法在 计算密集型任务（如深度学习、科学计算）中非常典型，其核心是平衡内存带宽与计算资源，最大化硬件利用率。

事实上，根据一开始代码运行出来的结果：Maximum number of registers per block: 64 K，每个block确实有64k寄存器，因此确实可以增大每个线程的计算量。

Deepseek版解析：

  1. 共享内存布局

    sa7 的布局：
    sa7 是一个64行×16列的共享内存数组，用于缓存矩阵A的当前分块数据。
    每个线程负责将A的4×1（4行1列）元素加载到sa7中。
    由于每个线程块有256个线程，每个线程处理A的一列中的4个元素，因此可以完整加载A的64行×16列分块。

    sb7 的布局：
    sb7 是一个16行×64列的共享内存数组，但以转置形式存储为64行×16列。
    每个线程负责将B的1×4（1行4列）元素加载到sb7中。
    由于每个线程块有256个线程，每个线程处理B的一列中的4个元素，因此可以完整加载B的16行×64列分块（转置后为64行×16列）。

  2. 大循环：加载数据到共享内存

    大循环的作用：
    沿K维度分块处理，每次处理16个K维度元素（KS_7=16）。
    在每个大循环迭代中，将A的64行×16列分块和B的16行×64列分块加载到共享内存。

    数据加载逻辑：
    每个线程根据tx（线程索引）确定自己负责加载A和B的哪部分数据。
    对于A：
    row_a = (tx & 15) << 2：每个线程负责A的4行。
    col_a = tx >> 4：每个线程负责A的某一列。
    加载A的4×1元素到sa7。

    对于B：
    row_b = (tx & 3) << 2：每个线程负责B的4列。
    col_b = tx >> 2：每个线程负责B的某一行。
    加载B的1×4元素并转置存储到sb7。

    更新全局内存指针：
    A += lda16：A按列主序，每次移动16列。
    B += 16：B按行主序，每次移动16行。
    为下一次大循环加载新的分块数据做准备。

    线程同步：
    __syncthreads()：确保所有线程完成共享内存的写入，避免数据竞争。

  3. 小循环：乘累加计算

    小循环的作用：
    对当前K分块（16次迭代）进行乘累加计算。
    每个线程从共享内存加载A的4×1元素和B的1×4元素，进行外积计算，累加到4×4的局部结果Cres中。

    计算逻辑：
    从sa7加载A的4×1元素（行向量）：

    1
    vload(Av, &sa7(row_a, inner_k_count));
    CPP

    从sb7加载B的1×4元素（列向量，已转置为行）：

    1
    vload(Bv, &sb7(col_c, inner_k_count));
    CPP

    进行外积计算并累加到Cres：

    1
    2
    3
    4
    vscal(Cres[0], Av, Bv.x); // Cres[0] += Av * Bv.x
    vscal(Cres[1], Av, Bv.y); // Cres[1] += Av * Bv.y
    vscal(Cres[2], Av, Bv.z); // Cres[2] += Av * Bv.z
    vscal(Cres[3], Av, Bv.w); // Cres[3] += Av * Bv.w
    CPP

    每个线程负责计算C中一个4×4的子块。

  4. 结果导出

    结果累加：
    将局部结果Cres与C的当前值结合（alpha * Cres + beta * C）。

    写回全局内存：
    使用float4将4×4结果写回C的对应位置，保持合并写入。

  5. 总结

    sa7 和 sb7 的布局：
    sa7 是64行×16列，缓存A的分块数据。
    sb7 是16行×64列，但以转置形式存储为64行×16列，缓存B的分块数据。

    大循环：
    根据tx确定索引，将A的64行×16列和B的16行×64列加载到共享内存。

    小循环：
    依次计算4行×1列的A和1行×4列的B的外积，得到4行×4列的结果，并累加到Cres。

    结果导出：
    将Cres的结果写回全局内存C。

1
2
3
4
5
6
7
int row_a = (tx&15)<<2, col_a = tx>>4;
    int row_b = (tx&3)<<2, col_b = tx>>2;
    int col_c = col_a<<2;
    int lda16 = lda<<4;
    A = &A((bx<<6),0);
    B = &B(0,(by<<6));
    C = &C((bx<<6),(by<<6));//the TB size is 64.
CPP

个人认为这段代码初看还是有一定难度的：

  • row_a：通过(tx & 15) << 2计算线程在A子块中的起始行。取tx的低4位（0-15），乘以4得到0,4,8,…,60。每个线程负责加载A的4行数据。
  • col_a：通过tx >> 4计算线程在A子块中的列。将tx右移4位，得到0-15的列索引。每个线程处理A的某一列。
  • row_b：通过(tx & 3) << 2计算线程在B子块中的起始行。取tx的低2位（0-3），乘以4得到0,4,8,12。每个线程加载B的4行数据。
  • col_b：通过tx >> 2计算线程在B子块中的列。将tx右移2位，得到0-63的列索引。每个线程处理B的某一列。

对于256个线程的线程块，就是依次处理A矩阵的（1-4行，第1列）、（5-8行，第1列）…（60-64行，第1列）然后（1-4行，第2列）和B矩阵的（1-4…12-16行，1-64列）。即把64行16列的A的子矩阵和16行64列的B的子矩阵导入共享内存。

然后A = &A((bx<<6),0); B = &B(0,(by<<6)); C = &C((bx<<6),(by<<6));就很好理解了，表示A往下移64行，B移动64列。

接下来的就是大循环，先把A和B中的子矩阵移入sa7和sb7，然后进入小循环，小循环中4*1的小矩阵和1*4的小矩阵相乘，形成4*4的矩阵存入Cres中，最后把Cres的值导入C，完成计算。

继续增大线程工作量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa8(i,j) sa8[((j)<<7) + (i)]
#define sb8(i,j) sb8[((j)<<7) + (i)]
#define MS_8 128
#define NS_8 128
#define KS_8 8
//v1 += v2 * s3, vector scaling
#define vscal(v1, v2, s3)\
    v1.x+=v2.x*s3;\
    v1.y+=v2.y*s3;\
    v1.z+=v2.z*s3;\
    v1.w+=v2.w*s3;
//v1 = alpha * v2 + beta * v3, simd fma
#define simd_axpby(v1, alpha, v2, beta, v3)\
    v1.x=alpha*v2.x+beta*v3.x;\
    v1.y=alpha*v2.y+beta*v3.y;\
    v1.z=alpha*v2.z+beta*v3.z;\
    v1.w=alpha*v2.w+beta*v3.w;
#define vload(v1,addr)\
    v1 = *((float4 *)(addr));
#define vstore(addr,v1)\
    *((float4 *)(addr)) = v1;
// cache blocking version, without register-level data re-use
// with memory coelascing on shared memory
// more workloads per thread. 8x8 micro kernel.
// adopt vetorized load/store
__global__  __launch_bounds__(256)
void mysgemm_v8(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int row_a = (tx&31)<<2, col_a = tx>>5;
    int row_b = (tx&1)<<2, col_b = tx>>1;
    int lda8 = lda<<3;
    int row_c = (tx&15)<<3, col_c = (tx>>4)<<3;
    A = &A((bx<<7),0);
    B = &B(0,(by<<7));
    C = &C((bx<<7),(by<<7));//the TB size is 128.
    __shared__ float sa8[1024];
    __shared__ float sb8[1024];
    float4 Av1, Av2, Bv1, Bv2, Cv[16], Cres[16];
    memset(Cres, 0, sizeof(Cres));//clear registers
    for (int k_count = 0; k_count<K; k_count+=KS_8){
        vload(Av1, &A(row_a,col_a))
        vload(Bv1, &B(row_b,col_b))
        ((float4 *)sa8)[tx] = Av1;
        sb8(col_b,row_b)=Bv1.x;
        sb8(col_b,row_b+1)=Bv1.y;
        sb8(col_b,row_b+2)=Bv1.z;
        sb8(col_b,row_b+3)=Bv1.w;
        A+=lda8;B+=8;
        __syncthreads();
        #pragma unroll
        for (int inner_k_count=0;inner_k_count<KS_8;inner_k_count++){
            vload(Av1, &sa8(row_c,inner_k_count))
            vload(Av2, &sa8(row_c+4,inner_k_count))
            vload(Bv1, &sb8(col_c,inner_k_count))
            vload(Bv2, &sb8(col_c+4,inner_k_count))
            vscal(Cres[0], Av1, Bv1.x)
            vscal(Cres[1], Av2, Bv1.x)
            vscal(Cres[2], Av1, Bv1.y)
            vscal(Cres[3], Av2, Bv1.y)
            vscal(Cres[4], Av1, Bv1.z)
            vscal(Cres[5], Av2, Bv1.z)
            vscal(Cres[6], Av1, Bv1.w)
            vscal(Cres[7], Av2, Bv1.w)
            vscal(Cres[8], Av1, Bv2.x)
            vscal(Cres[9], Av2, Bv2.x)
            vscal(Cres[10], Av1, Bv2.y)
            vscal(Cres[11], Av2, Bv2.y)
            vscal(Cres[12], Av1, Bv2.z)
            vscal(Cres[13], Av2, Bv2.z)
            vscal(Cres[14], Av1, Bv2.w)
            vscal(Cres[15], Av2, Bv2.w)
        }
        __syncthreads();
    }
    vload(Cv[0], &C(row_c,col_c))
    vload(Cv[1], &C(row_c+4,col_c))
    vload(Cv[2], &C(row_c,col_c+1))
    vload(Cv[3], &C(row_c+4,col_c+1))
    vload(Cv[4], &C(row_c,col_c+2))
    vload(Cv[5], &C(row_c+4,col_c+2))
    vload(Cv[6], &C(row_c,col_c+3))
    vload(Cv[7], &C(row_c+4,col_c+3))
    vload(Cv[8], &C(row_c,col_c+4))
    vload(Cv[9], &C(row_c+4,col_c+4))
    vload(Cv[10], &C(row_c,col_c+5))
    vload(Cv[11], &C(row_c+4,col_c+5))
    vload(Cv[12], &C(row_c,col_c+6))
    vload(Cv[13], &C(row_c+4,col_c+6))
    vload(Cv[14], &C(row_c,col_c+7))
    vload(Cv[15], &C(row_c+4,col_c+7))
    
    simd_axpby(Cres[0],alpha,Cres[0],beta,Cv[0])
    simd_axpby(Cres[1],alpha,Cres[1],beta,Cv[1])
    simd_axpby(Cres[2],alpha,Cres[2],beta,Cv[2])
    simd_axpby(Cres[3],alpha,Cres[3],beta,Cv[3])

    simd_axpby(Cres[4],alpha,Cres[4],beta,Cv[4])
    simd_axpby(Cres[5],alpha,Cres[5],beta,Cv[5])
    simd_axpby(Cres[6],alpha,Cres[6],beta,Cv[6])
    simd_axpby(Cres[7],alpha,Cres[7],beta,Cv[7])

    simd_axpby(Cres[8],alpha,Cres[8],beta,Cv[8])
    simd_axpby(Cres[9],alpha,Cres[9],beta,Cv[9])
    simd_axpby(Cres[10],alpha,Cres[10],beta,Cv[10])
    simd_axpby(Cres[11],alpha,Cres[11],beta,Cv[11])

    simd_axpby(Cres[12],alpha,Cres[12],beta,Cv[12])
    simd_axpby(Cres[13],alpha,Cres[13],beta,Cv[13])
    simd_axpby(Cres[14],alpha,Cres[14],beta,Cv[14])
    simd_axpby(Cres[15],alpha,Cres[15],beta,Cv[15])

    vstore(&C(row_c,col_c), Cres[0])
    vstore(&C(row_c+4,col_c), Cres[1])
    vstore(&C(row_c,col_c+1), Cres[2])
    vstore(&C(row_c+4,col_c+1), Cres[3])
    vstore(&C(row_c,col_c+2), Cres[4])
    vstore(&C(row_c+4,col_c+2), Cres[5])
    vstore(&C(row_c,col_c+3), Cres[6])
    vstore(&C(row_c+4,col_c+3), Cres[7])
    vstore(&C(row_c,col_c+4), Cres[8])
    vstore(&C(row_c+4,col_c+4), Cres[9])
    vstore(&C(row_c,col_c+5), Cres[10])
    vstore(&C(row_c+4,col_c+5), Cres[11])
    vstore(&C(row_c,col_c+6), Cres[12])
    vstore(&C(row_c+4,col_c+6), Cres[13])
    vstore(&C(row_c,col_c+7), Cres[14])
    vstore(&C(row_c+4,col_c+7), Cres[15])
}
CPP

本质和上一个写法没什么差别，不过是计算强度更强了：

  • 原共享内存的64*16变成了128*8
  • 原本是4*1的小矩阵和1*4的小矩阵相乘，现在变成了4*2和2*4的小矩阵相乘
  • 现在每个线程块负责128*128个元素的C的子矩阵

添加warp-level tiling和 warp-level parallelism

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#include<stdio.h>
#include<stdlib.h>
#define A(i,j) A[(i) + (j)*lda]
#define B(i,j) B[(i) + (j)*ldb]
#define C(i,j) C[(i) + (j)*ldc]
#define sa9(i,j) sa9[((j)<<7) + (i)]
#define sb9(i,j) sb9[((j)<<7) + (i)]
#define MS_9 128
#define NS_9 128
#define KS_9 8
//v1 += v2 * s3, vector scaling
#define vscal(v1, v2, s3)\
    v1.x+=v2.x*s3;\
    v1.y+=v2.y*s3;\
    v1.z+=v2.z*s3;\
    v1.w+=v2.w*s3;
//v1 = alpha * v2 + beta * v3, simd fma
#define simd_axpby(v1, alpha, v2, beta, v3)\
    v1.x=alpha*v2.x+beta*v3.x;\
    v1.y=alpha*v2.y+beta*v3.y;\
    v1.z=alpha*v2.z+beta*v3.z;\
    v1.w=alpha*v2.w+beta*v3.w;
#define vload(v1,addr)\
    v1 = *((float4 *)(addr));
#define vstore(addr,v1)\
    *((float4 *)(addr)) = v1;
// cache blocking version, without register-level data re-use
// with memory coelascing on shared memory
// more workloads per thread. 8x8 micro kernel.
// adopt vetorized load/store
__global__  __launch_bounds__(256)
void mysgemm_v9(int M, int N, int K, float alpha, float* A, float* B, float beta, float* C){
    int lda = M, ldb = K, ldc = M;
    int tx = threadIdx.x;
    int bx = blockIdx.x, by = blockIdx.y;
    int warp_id = tx>>5;
    int lane_id = tx&31;
    int warp_row = warp_id & 3, warp_col = warp_id >> 2;
    int row_w = lane_id&3, col_w = lane_id>>2;
    int row_b = (tx&1)<<2, col_b = tx>>1;
    int lda8 = lda<<3;
    int row_c = (warp_row<<5) + (row_w<<3), col_c = (warp_col<<6) + (col_w<<3);
    int row_a = (tx&31)<<2, col_a = tx>>5;
    A = &A((bx<<7),0);
    B = &B(0,(by<<7));
    C = &C((bx<<7),(by<<7));//the TB size is 128.
    __shared__ float sa9[1024];
    __shared__ float sb9[1024];
    float4 Av1, Av2, Bv1, Bv2, Cv[16], Cres[16];
    memset(Cres, 0, sizeof(Cres));//clear registers
    for (int k_count = 0; k_count<K; k_count+=KS_9){
        /*packing A and B into shared memory*/
        vload(Av1, &A(row_a,col_a))
        vload(Bv1, &B(row_b,col_b))
        ((float4 *)sa9)[tx] = Av1;
        sb9(col_b,row_b)=Bv1.x;
        sb9(col_b,row_b+1)=Bv1.y;
        sb9(col_b,row_b+2)=Bv1.z;
        sb9(col_b,row_b+3)=Bv1.w;
        A+=lda8;B+=8;
        __syncthreads();
        #pragma unroll
        for (int inner_k_count=0;inner_k_count<KS_9;inner_k_count++){
            vload(Av1, &sa9(row_c,inner_k_count))
            vload(Av2, &sa9(row_c+4,inner_k_count))
            vload(Bv1, &sb9(col_c,inner_k_count))
            vload(Bv2, &sb9(col_c+4,inner_k_count))
            vscal(Cres[0], Av1, Bv1.x)
            vscal(Cres[1], Av2, Bv1.x)
            vscal(Cres[2], Av1, Bv1.y)
            vscal(Cres[3], Av2, Bv1.y)
            vscal(Cres[4], Av1, Bv1.z)
            vscal(Cres[5], Av2, Bv1.z)
            vscal(Cres[6], Av1, Bv1.w)
            vscal(Cres[7], Av2, Bv1.w)
            vscal(Cres[8], Av1, Bv2.x)
            vscal(Cres[9], Av2, Bv2.x)
            vscal(Cres[10], Av1, Bv2.y)
            vscal(Cres[11], Av2, Bv2.y)
            vscal(Cres[12], Av1, Bv2.z)
            vscal(Cres[13], Av2, Bv2.z)
            vscal(Cres[14], Av1, Bv2.w)
            vscal(Cres[15], Av2, Bv2.w)
        }
        __syncthreads();
    }
    vload(Cv[0], &C(row_c,col_c))
    vload(Cv[1], &C(row_c+4,col_c))
    vload(Cv[2], &C(row_c,col_c+1))
    vload(Cv[3], &C(row_c+4,col_c+1))
    vload(Cv[4], &C(row_c,col_c+2))
    vload(Cv[5], &C(row_c+4,col_c+2))
    vload(Cv[6], &C(row_c,col_c+3))
    vload(Cv[7], &C(row_c+4,col_c+3))
    vload(Cv[8], &C(row_c,col_c+4))
    vload(Cv[9], &C(row_c+4,col_c+4))
    vload(Cv[10], &C(row_c,col_c+5))
    vload(Cv[11], &C(row_c+4,col_c+5))
    vload(Cv[12], &C(row_c,col_c+6))
    vload(Cv[13], &C(row_c+4,col_c+6))
    vload(Cv[14], &C(row_c,col_c+7))
    vload(Cv[15], &C(row_c+4,col_c+7))
    
    simd_axpby(Cres[0],alpha,Cres[0],beta,Cv[0])
    simd_axpby(Cres[1],alpha,Cres[1],beta,Cv[1])
    simd_axpby(Cres[2],alpha,Cres[2],beta,Cv[2])
    simd_axpby(Cres[3],alpha,Cres[3],beta,Cv[3])

    simd_axpby(Cres[4],alpha,Cres[4],beta,Cv[4])
    simd_axpby(Cres[5],alpha,Cres[5],beta,Cv[5])
    simd_axpby(Cres[6],alpha,Cres[6],beta,Cv[6])
    simd_axpby(Cres[7],alpha,Cres[7],beta,Cv[7])

    simd_axpby(Cres[8],alpha,Cres[8],beta,Cv[8])
    simd_axpby(Cres[9],alpha,Cres[9],beta,Cv[9])
    simd_axpby(Cres[10],alpha,Cres[10],beta,Cv[10])
    simd_axpby(Cres[11],alpha,Cres[11],beta,Cv[11])

    simd_axpby(Cres[12],alpha,Cres[12],beta,Cv[12])
    simd_axpby(Cres[13],alpha,Cres[13],beta,Cv[13])
    simd_axpby(Cres[14],alpha,Cres[14],beta,Cv[14])
    simd_axpby(Cres[15],alpha,Cres[15],beta,Cv[15])

    vstore(&C(row_c,col_c), Cres[0])
    vstore(&C(row_c+4,col_c), Cres[1])
    vstore(&C(row_c,col_c+1), Cres[2])
    vstore(&C(row_c+4,col_c+1), Cres[3])
    vstore(&C(row_c,col_c+2), Cres[4])
    vstore(&C(row_c+4,col_c+2), Cres[5])
    vstore(&C(row_c,col_c+3), Cres[6])
    vstore(&C(row_c+4,col_c+3), Cres[7])
    vstore(&C(row_c,col_c+4), Cres[8])
    vstore(&C(row_c+4,col_c+4), Cres[9])
    vstore(&C(row_c,col_c+5), Cres[10])
    vstore(&C(row_c+4,col_c+5), Cres[11])
    vstore(&C(row_c,col_c+6), Cres[12])
    vstore(&C(row_c+4,col_c+6), Cres[13])
    vstore(&C(row_c,col_c+7), Cres[14])
    vstore(&C(row_c+4,col_c+7), Cres[15])
}
CPP

前面的宏定义以及后面的共享内存赋值和计算部分和上一个核没有任何区别，主要的区别在于一开始的索引计算部分

kernel8的写法为：

1
2
3
4
5
6
7
8
9
10
int lda = M, ldb = K, ldc = M;
int tx = threadIdx.x;
int bx = blockIdx.x, by = blockIdx.y;
int row_a = (tx&31)<<2, col_a = tx>>5;
int row_b = (tx&1)<<2, col_b = tx>>1;
int lda8 = lda<<3;
int row_c = (tx&15)<<3, col_c = (tx>>4)<<3;
A = &A((bx<<7),0);
B = &B(0,(by<<7));
C = &C((bx<<7),(by<<7));//the TB size is 128.
CPP

而kernel9的写法为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
int lda = M, ldb = K, ldc = M;
int tx = threadIdx.x;
int bx = blockIdx.x, by = blockIdx.y;
int warp_id = tx>>5;
int lane_id = tx&31;
int warp_row = warp_id & 3, warp_col = warp_id >> 2;
int row_w = lane_id&3, col_w = lane_id>>2;
int row_b = (tx&1)<<2, col_b = tx>>1;
int lda8 = lda<<3;
int row_c = (warp_row<<5) + (row_w<<3), col_c = (warp_col<<6) + (col_w<<3);
int row_a = (tx&31)<<2, col_a = tx>>5;
A = &A((bx<<7),0);
B = &B(0,(by<<7));
C = &C((bx<<7),(by<<7));//the TB size is 128.
CPP

看了下知乎上的内存访问介绍,这篇文章讲bank conflict中的4byte和8byte两种情况讲的非常好

贴一段deepseek对这部分的解析：

这段代码实现了一个优化的矩阵乘法核函数（SGEMM），利用共享内存和线程束（warp）的合并访问特性，以提高计算效率。以下是代码的关键点解析：

  1. 线程与块的配置

    • 线程块（Block）：每个线程块包含256个线程（threadIdx.x范围0-255）。
      网格划分：网格中的块（blockIdx.x和blockIdx.y）对应输出矩阵C的划分，每个块负责计算C的一个128x128子块。
    • Warp划分：每个线程块分为8个warp（256线程/32线程每warp）。通过warp_id = tx >> 5计算warp索引，lane_id = tx & 31确定线程在warp内的位置。

  2. 内存访问优化

    • 共享内存（Shared Memory）：声明sa9和sb9各1024个float，用于缓存全局内存中的A和B矩阵子块，减少重复访问全局内存的开销。
    • 合并访问：同一warp内的线程访问共享内存时，通过{Mw, Nw} = {4xMr, 8xNr}的布局（如4x8的子块），使得内存访问可合并，提高带宽利用率。

  3. 索引计算

    • C矩阵的起始位置：每个线程块处理C的128x128子块，起始位置由块索引bx和by确定：
      C = &C[(bx << 7) + (by << 7) * ldc]; // 每个块处理128x128的C子块
    • 线程在C子块中的坐标：
      row_c和col_c通过warp内部分解（warp_row, warp_col, row_w, col_w）计算，每个warp负责C子块内的多个8x8区域。
      例如，row_c = (warp_row << 5) + (row_w << 3)，表示warp在行方向划分成4部分（每部分32行），每个线程进一步处理8行的数据。

  4. 全局内存到共享内存的加载

    • A和B的指针调整：全局内存中的A和B矩阵指针被偏移到当前块处理的起始位置：
      A = &A[(bx << 7)]; // A的起始行调整为bx128
      B = &B[(by << 7) * ldb]; // B的起始列调整为by      128
    • 共享内存布局：线程通过row_a, col_a, row_b, col_b等索引，将A和B的数据加载到共享内存。例如：
      row_a = (tx & 31) << 2表示每个线程加载A矩阵中连续的4行。
      col_b = tx >> 1结合row_b = (tx & 1) << 2，用于B矩阵的列优先布局，减少共享内存的bank冲突。

  5. 计算逻辑

    • 分块乘累加：虽然代码未展示完整计算循环，但典型SGEMM优化会分块加载A和B的子块到共享内存，进行局部乘累加。每个线程从共享内存读取数据，计算局部结果，并累加到C中。
    • Warp级并行：通过将线程划分为warp，并设计共享内存访问模式（如4x8子块），利用warp内32线程的SIMD并行性，同时合并内存访问。

  6. 性能优化

    • 减少Bank冲突：共享内存的布局设计（如row_b和col_b）确保同一warp内的线程访问不同bank，避免冲突。
    • 寄存器重用：通过循环展开和寄存器缓存中间结果，减少共享内存访问次数。
      总结

这段代码通过精细的线程划分、共享内存布局优化和索引计算，实现了高效的矩阵乘法。每个线程块处理128x128的C子块，利用共享内存缓存数据，减少全局内存访问，并通过warp级并行和合并访问提升性能。设计核心在于平衡线程工作量、内存访问模式与并行度，充分发挥GPU的计算能力。

gemm学习笔记
http://zzsy.me/2025/02/23/gemm学习笔记/
作者
yuanyuan
发布于
2025年2月23日
更新于
2025年2月25日
许可协议