From d13b493d40f42005782fb2c4c1c6841f4fa35bd1 Mon Sep 17 00:00:00 2001 From: yukun-hh Date: Sun, 27 Sep 2026 14:15:21 +0800 Subject: [PATCH] Dense GEMM --- 2(Gemm)/main1.cu | 60 ++++++++++++++++++++++++++++ 2(Gemm)/main2.cu | 102 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 162 insertions(+) create mode 100644 2(Gemm)/main1.cu create mode 100644 2(Gemm)/main2.cu diff --git a/2(Gemm)/main1.cu b/2(Gemm)/main1.cu new file mode 100644 index 0000000..fd1ad7b --- /dev/null +++ b/2(Gemm)/main1.cu @@ -0,0 +1,60 @@ +#include +#include + +#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<<>>(A, B, C, M, N, K); +} \ No newline at end of file diff --git a/2(Gemm)/main2.cu b/2(Gemm)/main2.cu new file mode 100644 index 0000000..caaf4ab --- /dev/null +++ b/2(Gemm)/main2.cu @@ -0,0 +1,102 @@ +#include +#include +#include + +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 a_frag; + wmma::fragment b_frag; + wmma::fragment 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<<>>(A, B, C, M, N, K); +} \ No newline at end of file