data: cat 3 SGEMM repos — siboehm, wangzyon, edtallison (full clone, no --depth)

Sources: siboehm/SGEMM_CUDA → upstream_ref/sgemm_siboehm/ (25 files) wangzyon/NVIDIA_SGEMM_PRACTICE → upstream_ref/nvidia_sgemm_practice/ (23 files, filled gaps) edtallison/sgemm-cuda → upstream_ref/sgemm_edtallison/ (41 files) All files cat'd one by one from git clone (no --depth). These are the 3 public SGEMM repos that can compile on CUDA 10.2 + CoreX ivcore10. Key files for BI-V100 porting: kernel 10 (warp tiling) — already proven on device with WARPSIZE=64 kernel 11/12 (double buffering) — next optimization target sgemm.cu + runner.cu — complete build+benchmark harness CMakeLists.txt — build system reference
2026-08-15 06:58:07 +00:00
parent 854fb93a8e
commit 284804ac53
65 changed files with 6963 additions and 0 deletions
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_1.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_1.cuh
@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+__global__ __launch_bounds__(1024) void
+mysgemm_v1(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+
+    int gx = blockIdx.x * blockDim.x + threadIdx.x; // 全局x
+    int gy = blockIdx.y * blockDim.y + threadIdx.y; // 全局y
+
+    float tmp = 0.;
+    for (int i = 0; i < K; i++) {
+        tmp += A[gy * K + i] * B[i * N + gx]; // 两次全局内存访问和一次FMA（累加乘）
+    }
+    C[gy * N + gx] = alpha * tmp + beta * C[gy * N + gx];
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_2.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_2.cuh
@@ -0,0 +1,45 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BLOCK_SIZE>
+__global__ void mysgemm_v2(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int BM = BLOCK_SIZE;
+    const int BN = BLOCK_SIZE;
+    const int BK = BLOCK_SIZE;
+
+    int tx = threadIdx.x % BN;
+    int ty = threadIdx.x / BN;
+
+    // 申请共享内存空间
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    float tmp = 0.;
+    for (int k = 0; k < K; k += BK) {
+        // 缓存A_tile和B_tile
+        As[ty * BK + tx] = A[ty * K + tx];
+        Bs[ty * BN + tx] = B[ty * N + tx];
+        // 同步所有线程缓存完成
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        for (int i = 0; i < BK; i++) {
+            tmp += As[ty * BK + i] * Bs[i * BN + tx];
+        }
+        // FMA计算需要读取缓存数据，在新一轮写入缓存前进行同步，确保所有线程计算完成
+        __syncthreads();
+    }
+    C[ty * N + tx] = alpha * tmp + beta * C[ty * N + tx];
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_3.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_3.cuh
@@ -0,0 +1,71 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM>
+__global__ void mysgemm_v3(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+    int thread_num = BM * BN / TM; // 一个线程负责block中计算TM个元素
+
+    int tx = threadIdx.x % BN;
+    int ty = threadIdx.x / BN * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM + 1] = {0.}; // 每个线程负责TM个元素，则需要申请TM个寄存器保存累加值，额外的一个寄存器用于缓存；
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            tmp[TM] = Bs[tx + i * BN]; // 额外的一个寄存器，避免反复从共享内存中读取Bs[tx + i * BN]
+            #pragma unroll  // 循环展开，增加指令并行度
+            for (int j = 0; j < TM; j++) {
+                tmp[j] += As[(ty + j) * BK + i] * tmp[TM];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        C[(ty + j) * N + tx] = alpha * tmp[j] + beta * C[(ty + j) * N + tx];
+    }
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_4.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_4.cuh
@@ -0,0 +1,76 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v4(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    int block_row_thread = BN / TN;
+    int block_col_thread = BM / TM;
+    int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            #pragma unroll  // 循环展开，增加指令并行度
+            for (int j = 0; j < TM; j++) {
+                for (int l = 0; l < TN; l++)
+                    tmp[j][l] += As[(ty + j) * BK + i] * Bs[tx + l + i * BN];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        for (int l = 0; l < TN; l++)
+            C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l];
+    }
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_5.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_5.cuh
@@ -0,0 +1,88 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v5(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    int block_row_thread = BN / TN;
+    int block_col_thread = BM / TM;
+    int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BM * BK];
+    __shared__ float Bs[BK * BN];
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    /*
+    当前线程负责搬运全局内存中第a_tile_row行，第a_tile_col列元素至共享内存第a_tile_row行，第a_tile_col列
+    a_tile_stride表示block中线程可搬运a_tile_stride行至共享内存；
+
+    若BM=64,BK=8,thread_num=512,则a_tile_stride=64,a_tile_stride=BM，表示每个线程搬运一轮即可完成所需元素的搬运;
+    若BM=128,BK=8,thread_num=512,则a_tile_stride=64,表示每个线程搬运两轮即可完成所需元素的搬运;
+    */
+    int a_tile_row = threadIdx.x / BK;
+    int a_tile_col = threadIdx.x % BK;
+    int a_tile_stride = thread_num / BK;
+
+    int b_tile_row = threadIdx.x / BN;
+    int b_tile_col = threadIdx.x % BN;
+    int b_tile_stride = thread_num / BN;
+
+    float tmp[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+    float a_frag[TM] = {0.};
+    float b_frag[TN] = {0.};
+
+    #pragma unroll
+    for (int k = 0; k < K; k += BK) {
+        #pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            As[(a_tile_row + i) * BK + a_tile_col] = A[(a_tile_row + i) * K + a_tile_col];
+        }
+        #pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            Bs[(b_tile_row + i) * BN + b_tile_col] = B[(b_tile_row + i) * N + b_tile_col];
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+        #pragma unroll
+        for (int i = 0; i < BK; i++) {
+            #pragma unroll
+            for (int j = 0; j < TM; j++) {
+                a_frag[j] = As[(ty + j) * BK + i];
+            }
+            #pragma unroll
+            for (int l = 0; l < TN; l++) {
+                b_frag[l] = Bs[tx + l + i * BN];
+            }
+            #pragma unroll
+            for (int j = 0; j < TM; j++) {
+                #pragma unroll
+                for (int l = 0; l < TN; l++)
+                    tmp[j][l] += a_frag[j] * b_frag[l];
+            }
+        }
+        __syncthreads();
+    }
+    #pragma unroll
+    for (int j = 0; j < TM; j++) {
+        for (int l = 0; l < TN; l++)
+            C[(ty + j) * N + tx + l] = alpha * tmp[j][l] + beta * C[(ty + j) * N + tx + l];
+    }
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_6.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_6.cuh
@@ -0,0 +1,110 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v6(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int block_row_thread = BN / TN;
+    const int block_col_thread = BM / TM;
+    const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    // 当前线程对应thread tile的左上角元素在block中的位置
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[BK * BM];
+    __shared__ float Bs[BK * BN];
+
+
+    const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮
+    const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮
+
+    int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运
+    int a_tile_col = threadIdx.x % (BK / 4) * 4;
+    int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行
+
+    int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运
+    int b_tile_col = threadIdx.x % (BN / 4) * 4;
+    int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行
+
+    float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+
+    // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小
+    float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+
+    float a_frag[TM];  // 缓存As共享内存
+    float b_frag[TN];  // 缓存Bs共享内存
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+#pragma unroll
+    for (int k = 0; k < K; k += BK) {
+#pragma unroll
+        for (int i = 0; i < BM; i += a_tile_stride) {
+            int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+            FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                    FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]);
+            // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取
+            As[OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+            As[OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+            As[OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+            As[OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+        }
+#pragma unroll
+        for (int i = 0; i < BK; i += b_tile_stride) {
+            FETCH_FLOAT4(Bs[OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                    FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置
+        }
+        __syncthreads();
+        A += BK;
+        B += BK * N;
+#pragma unroll
+        for (int i = 0; i < BK; i++) {
+#pragma unroll
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[m]) = FETCH_FLOAT4(As[OFFSET(i, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[n]) = FETCH_FLOAT4(Bs[OFFSET(i, tx + n, BN)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int m = 0; m < TM; m++) {
+#pragma unroll
+                for (int n = 0; n < TN; n++) {
+                    accum[m][n] += a_frag[m] * b_frag[n];
+                }
+            }
+        }
+        __syncthreads();
+    }
+#pragma unroll
+    for (int m = 0; m < TM; m++) {
+#pragma unroll
+        for (int n = 0; n < TN; n += 4) {
+            float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]);
+            //float4 atmp = FETCH_FLOAT4(accum[m][n]);
+            ctmp.x = alpha * accum[m][n] + beta * ctmp.x;
+            ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y;
+            ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z;
+            ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w;
+            FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp;
+        }
+    }
+}
--- a/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_7.cuh
+++ b/upstream_ref/nvidia_sgemm_practice/src/kernel/kernel_7.cuh
@@ -0,0 +1,180 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define OFFSET(row, col, ld) ((row)*(ld)+(col))
+#define FETCH_FLOAT4(pointer) (reinterpret_cast<float4*>(&(pointer))[0])
+
+template<const int BM,
+        const int BN,
+        const int BK,
+        const int TM,
+        const int TN>
+__global__ void mysgemm_v7(int M, int N, int K, float alpha, float *A, float *B, float beta, float *C) {
+    int bx = blockIdx.x;
+    int by = blockIdx.y;
+
+    const int block_row_thread = BN / TN;
+    const int block_col_thread = BM / TM;
+    const int thread_num = block_row_thread * block_col_thread; // 一个线程负责计算block中TM*TN个元素
+
+    // 当前线程对应thread tile的左上角元素在block中的位置
+    int tx = (threadIdx.x % block_row_thread) * TN;
+    int ty = (threadIdx.x / block_row_thread) * TM;
+
+    __shared__ float As[2][BK * BM]; // 增加一倍共享内存大小用于缓存
+    __shared__ float Bs[2][BK * BN];
+
+
+    const int ldg_a_num = BK * BM / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至As需要所有线程搬运ldg_a_num轮
+    const int ldg_b_num = BK * BN / thread_num / 4; // 每个线程搬运4个浮点数，完成搬运至Bs需要所有线程搬运ldg_b_num轮
+
+    int a_tile_row = threadIdx.x / (BK / 4); // 每行4个字节作为一个内存块，当前线程负责第a_tile_row行的第a_tile_col个内存块的搬运
+    int a_tile_col = threadIdx.x % (BK / 4) * 4;
+    int a_tile_stride = BM / ldg_a_num; // 一共BM行，搬运ldg_a_num轮，每论搬运a_tile_stride行
+
+    int b_tile_row = threadIdx.x / (BN / 4); // 每行4个字节作为一个内存块，当前线程负责第b_tile_row行的第b_tile_col个内存块的搬运
+    int b_tile_col = threadIdx.x % (BN / 4) * 4;
+    int b_tile_stride = BK / ldg_b_num; // 一共BK行，搬运ldg_b_num轮，每论搬运b_tile_stride行
+
+    float accum[TM][TN] = {0.}; // 每个线程负责TM*TN个元素，则需要申请TM*TN个寄存器保存累加值，额外的一个寄存器用于缓存；
+
+    // 计算ldg_a_num的所有参数必须全部是const，否则不能用来申明数组大小
+    float ldg_a_reg[4 * ldg_a_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+    float ldg_b_reg[4 * ldg_b_num] = {0.}; // 每个线程搬运ldg_a_num轮，寄存器缓存ldg_a_num个float4元素，用于转置As矩阵
+
+    float a_frag[2][TM];  // 缓存As共享内存,增加一倍寄存器大小用于缓存
+    float b_frag[2][TN];  // 缓存Bs共享内存,增加一倍寄存器大小用于缓存
+
+    // 移动到当前block
+    A = &A[by * BM * K];
+    B = &B[bx * BN];
+    C = &C[by * BM * N + bx * BN];
+
+    // first global to shared
+#pragma unroll
+    for (int i = 0; i < BM; i += a_tile_stride) {
+        int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+        FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                FETCH_FLOAT4(A[OFFSET(a_tile_row + i, a_tile_col, K)]);
+        // As转置存，其中ldg_a_reg做中间缓存，目的是读取时可以按FLOAT4读取
+        As[0][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+        As[0][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+        As[0][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+        As[0][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+    }
+#pragma unroll
+    for (int i = 0; i < BK; i += b_tile_stride) {
+        FETCH_FLOAT4(Bs[0][OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                FETCH_FLOAT4(B[OFFSET(b_tile_row + i, b_tile_col, N)]); // 不需要转置
+    }
+    __syncthreads();
+
+    // first shared to frag
+#pragma unroll
+    for (int m = 0; m < TM; m += 4) {
+        FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4(As[0][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile
+    }
+#pragma unroll
+    for (int n = 0; n < TN; n += 4) {
+        FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4(Bs[0][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile
+    }
+
+
+    int write_index = 1;
+    int load_index;
+    int k = 0;
+    do {
+        k += BK;
+        // load global to reg
+        if (k < K) {
+#pragma unroll
+            for (int i = 0; i < BM; i += a_tile_stride) {
+                int ldg_index = i / a_tile_stride * 4;  // 第ldg_index轮
+                FETCH_FLOAT4(ldg_a_reg[ldg_index]) =
+                        FETCH_FLOAT4(A[OFFSET(a_tile_row + i, k + a_tile_col, K)]);
+            }
+#pragma unroll
+            for (int i = 0; i < BK; i += b_tile_stride) {
+                int ldg_index = i / b_tile_stride * 4;  // 第ldg_index轮
+                FETCH_FLOAT4(ldg_b_reg[ldg_index]) =
+                        FETCH_FLOAT4(B[OFFSET(k + b_tile_row + i, b_tile_col, N)]);
+            }
+        }
+
+        load_index = write_index ^ 1;
+#pragma unroll
+        for (int bk = 0; bk < BK - 1; bk++) {
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[(bk + 1) % 2][m]) = FETCH_FLOAT4(
+                        As[load_index][OFFSET(bk + 1, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[(bk + 1) % 2][n]) = FETCH_FLOAT4(
+                        Bs[load_index][OFFSET(bk + 1, tx + n, BN)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int m = 0; m < TM; m++) {
+                for (int n = 0; n < TN; n++) {
+                    accum[m][n] += a_frag[bk % 2][m] * b_frag[bk % 2][n];
+                }
+            }
+        }
+        if (k < K) {
+#pragma unroll
+            for (int i = 0; i < BM; i += a_tile_stride) {
+                int ldg_index = i / a_tile_stride * 4;
+                As[write_index][OFFSET(a_tile_col, i + a_tile_row, BM)] = ldg_a_reg[ldg_index];
+                As[write_index][OFFSET(a_tile_col + 1, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 1];
+                As[write_index][OFFSET(a_tile_col + 2, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 2];
+                As[write_index][OFFSET(a_tile_col + 3, i + a_tile_row, BM)] = ldg_a_reg[ldg_index + 3];
+            }
+#pragma unroll
+            for (int i = 0; i < BK; i += b_tile_stride) {
+                int ldg_index = i / b_tile_stride * 4;
+                FETCH_FLOAT4(Bs[write_index][OFFSET(b_tile_row + i, b_tile_col, BN)]) =
+                        FETCH_FLOAT4(ldg_b_reg[ldg_index]);
+            }
+            __syncthreads();
+#pragma unroll
+            for (int m = 0; m < TM; m += 4) {
+                FETCH_FLOAT4(a_frag[0][m]) = FETCH_FLOAT4(
+                        As[write_index][OFFSET(0, ty + m, BM)]); // 偏移到当前thread tile
+            }
+#pragma unroll
+            for (int n = 0; n < TN; n += 4) {
+                FETCH_FLOAT4(b_frag[0][n]) = FETCH_FLOAT4(
+                        Bs[write_index][OFFSET(0, tx + n, BN)]); // 偏移到当前thread tile
+            }
+
+            write_index ^= 1;
+        }
+#pragma unroll
+        for (int m = 0; m < TM; m++) {
+#pragma unroll
+            for (int n = 0; n < TN; n++) {
+                accum[m][n] += a_frag[(BK - 1) % 2][m] * b_frag[(BK - 1) % 2][n];
+            }
+        }
+
+
+    } while (k < K);
+    
+    // C = alpha*AB+C
+#pragma unroll
+    for (int m = 0; m < TM; m++) {
+#pragma unroll
+        for (int n = 0; n < TN; n += 4) {
+            float4 ctmp = FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]);
+            ctmp.x = alpha * accum[m][n] + beta * ctmp.x;
+            ctmp.y = alpha * accum[m][n + 1] + beta * ctmp.y;
+            ctmp.z = alpha * accum[m][n + 2] + beta * ctmp.z;
+            ctmp.w = alpha * accum[m][n + 3] + beta * ctmp.w;
+            FETCH_FLOAT4(C[OFFSET(ty + m, tx + n, N)]) = ctmp;
+        }
+    }
+}