Dense GEMM

2026-09-27 14:15:21 +08:00
parent 71fd142b03
commit d13b493d40
2 changed files with 162 additions and 0 deletions
@@ -0,0 +1,60 @@
+#include <stdint.h>
+#include <cuda_bf16.h>
+
+#define TILE 16
+__global__ void MatMulShareBf16(const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+                                int64_t M, int64_t N, int64_t K) {
+    __shared__ __nv_bfloat16 As[TILE][TILE];
+    __shared__ __nv_bfloat16 Bs[TILE][TILE];
+    
+    int tx = threadIdx.x, ty = threadIdx.y;
+    int row = blockIdx.y * TILE + ty; // 对应输出 C 的 M 维度 (i)
+    int col = blockIdx.x * TILE + tx; // 对应输出 C 的 N 维度 (j)
+    
+    float sum = 0.0f; // 使用 FP32 进行累加，保证计算精度
+    
+    int num = (K + TILE - 1) / TILE;
+    
+    for (int t = 0; t < num; t++) {
+        int aCol = tx + t * TILE;
+        int bK   = tx + t * TILE;
+        int bN   = blockIdx.x * TILE + ty;   // 注意这里用 ty 作为 N 方向
+
+        // A: A[row][k]
+        As[ty][tx] = (row < M && aCol < K)
+                        ? A[row * K + aCol]
+                                : __float2bfloat16(0.0f);
+
+        // B: B 是 (N,K)，这里存成 Bs[K][N]
+        Bs[tx][ty] = (bN < N && bK < K)
+                                ? B[bN * K + bK]
+                                : __float2bfloat16(0.0f);
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE; ++k) {
+            float a_val = __bfloat162float(As[ty][k]);
+            float b_val = __bfloat162float(Bs[k][tx]); // Bs[K][N]
+            sum += a_val * b_val;
+        }
+        
+        __syncthreads();
+    }
+    
+    // 将 FP32 结果转换回 bf16 写回全局内存
+    if (row < M && col < N) {
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+extern "C" void run_kernel(
+    const __nv_bfloat16* A,
+    const __nv_bfloat16* B,
+    __nv_bfloat16* C,
+    int64_t M,
+    int64_t N,
+    int64_t K
+){
+    dim3 block16(TILE, TILE);
+    dim3 grid16((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);
+    MatMulShareBf16<<<grid16, block16>>>(A, B, C, M, N, K);
+}
@@ -0,0 +1,102 @@
+#include <cuda_bf16.h>
+#include <mma.h>
+#include <cstdint>
+
+using namespace nvcuda;
+
+constexpr int BM = 32;        // block tile M
+constexpr int BN = 32;        // block tile N
+constexpr int BK = 128;        // block tile K
+constexpr int PAD = 4;        // shared memory padding，缓解 bank 冲突
+
+constexpr int WARPS_M = BM / 16;   // 4
+constexpr int WARPS_N = BN / 16;   // 4
+constexpr int NUM_WARPS = WARPS_M * WARPS_N;   // 16
+constexpr int THREADS   = NUM_WARPS * 32;      // 512
+
+__global__ void wmma_bf16_gemm_abt(const __nv_bfloat16* __restrict__ A,
+                                   const __nv_bfloat16* __restrict__ B,
+                                   __nv_bfloat16* __restrict__ C,
+                                   int M, int N, int K) {
+    const int tid     = threadIdx.x;
+    const int warp_id = tid / 32;
+    const int lane    = tid % 32;
+    const int warp_m  = warp_id / WARPS_N;
+    const int warp_n  = warp_id % WARPS_N;
+
+    const int block_m = blockIdx.y * BM;
+    const int block_n = blockIdx.x * BN;
+
+    __shared__ __nv_bfloat16 A_smem[BM][BK + PAD];
+    __shared__ __nv_bfloat16 B_smem[BN][BK + PAD];   // 按 N x K 行主序存
+    __shared__ float         C_smem[NUM_WARPS][16][16];
+
+    wmma::fragment<wmma::matrix_a, 16, 16, 16, __nv_bfloat16, wmma::row_major> a_frag;
+    wmma::fragment<wmma::matrix_b, 16, 16, 16, __nv_bfloat16, wmma::col_major> b_frag;
+    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
+
+    wmma::fill_fragment(c_frag, 0.0f);
+
+    for (int k0 = 0; k0 < K; k0 += BK) {
+        // -------- 协作加载 A tile: BM x BK --------
+        for (int i = tid; i < BM * BK; i += THREADS) {
+            int r = i / BK;
+            int c = i % BK;
+            int gr = block_m + r;
+            int gc = k0 + c;
+            A_smem[r][c] = (gr < M && gc < K) ? A[gr * K + gc]
+                                              : __float2bfloat16(0.0f);
+        }
+        // -------- 协作加载 B tile: BN x BK --------
+        for (int i = tid; i < BN * BK; i += THREADS) {
+            int r = i / BK;
+            int c = i % BK;
+            int gr = block_n + r;
+            int gc = k0 + c;
+            B_smem[r][c] = (gr < N && gc < K) ? B[gr * K + gc]
+                                              : __float2bfloat16(0.0f);
+        }
+        __syncthreads();
+
+       for (int kk = 0; kk < BK; kk += 16) {
+            // A 片段：从 shared 的 kk 列开始，取 16 列
+            wmma::load_matrix_sync(a_frag,
+                               &A_smem[warp_m * 16][kk],
+                               BK + PAD);
+
+            // B 片段：从 shared 的 kk 列开始，取 16 列，col_major 实现 B^T
+            wmma::load_matrix_sync(b_frag,
+                               &B_smem[warp_n * 16][kk],
+                               BK + PAD);
+
+            wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
+        }
+        __syncthreads();   // 下一轮覆盖 shared memory 前同步
+    }
+
+    // -------- 结果写回 --------
+    wmma::store_matrix_sync(&C_smem[warp_id][0][0], c_frag, 16,
+                            wmma::mem_row_major);
+
+    const int tile_m = block_m + warp_m * 16;
+    const int tile_n = block_n + warp_n * 16;
+
+    for (int i = lane; i < 16 * 16; i += 32) {
+        int r = i / 16;
+        int c = i % 16;
+        int gr = tile_m + r;
+        int gc = tile_n + c;
+        if (gr < M && gc < N) {
+            C[gr * N + gc] = __float2bfloat16(C_smem[warp_id][r][c]);
+        }
+    }
+}
+
+extern "C" void run_kernel(const __nv_bfloat16* A,
+                           const __nv_bfloat16* B,
+                           __nv_bfloat16* C,
+                           int M, int N, int K) {
+    dim3 block(THREADS);
+    dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
+    wmma_bf16_gemm_abt<<<grid, block>>>(A, B, C, M, N, K);
+}