Sparse Matrix-Dense Matrix Multiplication

原始题目：LeetGPU - Sparse Matrix-Dense Matrix Multiplication

题目描述

编写一个 GPU 程序，将稀疏矩阵 $A$（$M \times N$）与稠密矩阵 $B$（$N \times K$）相乘，产生稠密输出矩阵 $C$（$M \times K$）。所有矩阵以行优先顺序存储，使用 32 位浮点数。矩阵 $A$ 约 60%–70% 稀疏（即 60%–70% 的元素为零），$nnz$ 为 $A$ 中非零元素的数量。

数学上：

$$C_{ij} = \sum_{k=0}^{N-1} A_{ik} \cdot B_{kj}, \quad i = 0,\dots,M-1,\ j = 0,\dots,K-1$$

实现要求

• 只允许使用 GPU 原生功能（不允许使用外部库）。
• solve 函数签名必须保持不变。
• 最终结果必须存储在矩阵 $C$ 中。

示例

Input:  A (3×4): [[2,0,0,1],[0,3,0,0],[0,0,4,0]]
        B (4×2): [[1,2],[3,4],[5,6],[7,8]]
Output: C (3×2): [[9,12],[9,12],[20,24]]

约束条件

• $1 \le M, N, K \le 8{,}192$。
• $A$ 和 $B$ 中的值范围 $[-10, 10]$，FP32。
• $A$ 约 60%–70% 稀疏。
• 性能测试在 $M = 4{,}096,\ N = 2{,}048,\ K = 512$ 的规模下进行。

解题思路

SpMM 可以看作对 $A$ 的每个非零元素 $A_{ik}$，将 $B$ 的第 $k$ 行乘以该值加到 $C$ 的第 $i$ 行。如果每行非零元素数量变化大，使用 CSR 格式并让每个 warp 处理一行（或动态分配）更高效。对于高稀疏度的场景（60-70%），可以考虑直接转为稠密格式用 GEMM 代替。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void spmm_kernel(const float* values, const int* col_idx, const int* row_ptr,
    const float* B, float* C, int M, int K) {
    int row = blockIdx.x;
    int col = threadIdx.x + blockIdx.y * blockDim.x;
    if (row < M && col < K) {
        float sum = 0.0f;
        for (int j = row_ptr[row]; j < row_ptr[row+1]; j++)
            sum += values[j] * B[col_idx[j] * K + col];
        C[row * K + col] = sum;
    }
}
extern "C" void solve(const float* values, const int* col_idx, const int* row_ptr,
    const float* B, float* C, int M, int N, int K) {
    dim3 threads(32, 8), blocks(M, (K+7)/8);
    spmm_kernel<<<blocks, threads>>>(values, col_idx, row_ptr, B, C, M, K);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def spmm_kernel(values_ptr, col_idx_ptr, row_ptr_ptr, B_ptr, C_ptr,
    M: tl.constexpr, K: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    col = tl.program_id(1)*BLOCK + tl.arange(0, BLOCK)
    mask = col < K
    start = tl.load(row_ptr_ptr + row); end = tl.load(row_ptr_ptr + row + 1)
    acc = tl.zeros((BLOCK,), tl.float32)
    for j in range(end - start):
        vcol = tl.load(col_idx_ptr + start + j)
        val = tl.load(values_ptr + start + j)
        acc += val * tl.load(B_ptr + vcol * K + col, mask=mask)
    tl.store(C_ptr + row * K + col, acc, mask=mask)