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 代替。欢迎在 GitHub Discussions 分享你的解法。