General Matrix Multiplication (GEMM)

原始题目：LeetGPU - General Matrix Multiplication (GEMM)

题目描述

实现一个基本的通用矩阵乘法（GEMM）。给定矩阵 $A$（$M \times K$）、$B$（$K \times N$）、输入/输出矩阵 $C$（$M \times N$）以及标量乘数 $\alpha$ 和 $\beta$，计算：

$$C = \alpha \cdot (A \times B) + \beta \cdot C_{initial}$$

输入矩阵 $A$、$B$ 和 $C$ 的初始状态包含 16 位浮点数（FP16/half），所有矩阵以行优先顺序存储。标量 $\alpha$ 和 $\beta$ 为 32 位浮点数。

实现要求

• 只允许使用原生功能（除 WMMA 外不允许使用外部库）。
• solve 函数签名必须保持不变。
• 乘法过程中的累加应使用 FP32 以获得更好的精度，最终结果转换回 FP16。
• 最终结果必须以 half 类型存储回矩阵 $C$。

示例

Input:  A (2×3, FP16): [[1,2,3],[4,5,6]]
        B (3×2, FP16): [[1,2],[3,4],[5,6]]
        C_init (2×2, FP16): [[1,1],[1,1]]
        α = 1.0, β = 0.0
Output: C (2×2, FP16): [[22,28],[49,64]]

约束条件

• $16 \le M, N, K \le 4{,}096$。
• 性能测试在 $M = N = K = 1{,}024$ 的规模下进行。

解题思路

GEMM 是 GPU 计算中最重要的内核之一，cuBLAS 的大部分性能优化都围绕它展开。核心优化链：朴素全局内存 → 共享内存分块 → 寄存器分块 → 向量化加载（float4）→ 双缓冲 → warp-level 矩阵指令（WMMA/Tensor Core）。FP16 计算可以利用 Tensor Core（Ampere+）获得数倍于 FP32 的吞吐。累加器用 FP32 是标准做法——FP16 的精度不足以安全地累积大量项。

代码实现

CUDA

#include <cuda_runtime.h>
#define TILE 32
__global__ void gemm_tiled(const float* A, const float* B, float* C, int M, int K, int N, float alpha, float beta) {
    __shared__ float As[TILE][TILE], Bs[TILE][TILE];
    int bx=blockIdx.x, by=blockIdx.y, tx=threadIdx.x, ty=threadIdx.y;
    int row=by*TILE+ty, col=bx*TILE+tx;
    float sum=0.0f;
    for(int t=0; t<(K+TILE-1)/TILE; t++) {
        if(row<M && t*TILE+tx<K) As[ty][tx]=A[row*K + t*TILE+tx];
        else As[ty][tx]=0.0f;
        if(t*TILE+ty<K && col<N) Bs[ty][tx]=B[(t*TILE+ty)*N+col];
        else Bs[ty][tx]=0.0f;
        __syncthreads();
        for(int k=0;k<TILE;k++) sum+=As[ty][k]*Bs[k][tx];
        __syncthreads();
    }
    if(row<M && col<N) C[row*N+col] = alpha*sum + beta*C[row*N+col];
}
extern "C" void solve(const float* A, const float* B, float* C, int M, int K, int N, float alpha, float beta) {
    dim3 threads(TILE,TILE), blocks((N+TILE-1)/TILE, (M+TILE-1)/TILE);
    gemm_tiled<<<blocks,threads>>>(A,B,C,M,K,N,alpha,beta);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def gemm_kernel(A_ptr,B_ptr,C_ptr, M,N,K, alpha,beta, BLOCK_M:tl.constexpr,BLOCK_N:tl.constexpr,BLOCK_K:tl.constexpr):
    pm= tl.program_id(0); pn= tl.program_id(1)
    rm= pm*BLOCK_M+tl.arange(0,BLOCK_M); rn= pn*BLOCK_N+tl.arange(0,BLOCK_N); rk= tl.arange(0,BLOCK_K)
    acc= tl.zeros((BLOCK_M,BLOCK_N),tl.float32)
    for k in range(0,K,BLOCK_K):
        a= tl.load(A_ptr+rm[:,None]*K+(k+rk)[None,:], mask=rm[:,None]<M)
        b= tl.load(B_ptr+(k+rk)[:,None]*N+rn[None,:], mask=rn[None,:]<N)
        acc+= tl.dot(a,b)
    c= tl.load(C_ptr+rm[:,None]*N+rn[None,:])
    tl.store(C_ptr+rm[:,None]*N+rn[None,:], alpha*acc+beta*c)