Sparse Matrix-Vector Multiplication

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

题目描述

编写一个 GPU 程序，执行稀疏矩阵-向量乘法（SpMV）。给定一个 $M \times N$ 的稀疏矩阵 $A$ 和一个长度为 $N$ 的稠密向量 $x$，计算乘积向量 $y = A \times x$（长度为 $M$）。矩阵 $A$ 以行优先顺序存储，$nnz$ 为 $A$ 中非零元素的数量。

数学上：

$$y_i = \sum_{j=0}^{N-1} A_{ij} \cdot x_j,\quad i = 0, 1, \dots, M-1$$

矩阵 $A$ 约 60%–70% 的元素为零。

实现要求

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

示例

Input:  A (3×4): [[5,0,0,1], [0,2,3,0], [0,0,0,4]]
        x: [1, 2, 3, 4]
Output: y: [9, 13, 16]

约束条件

• $1 \le M, N \le 10{,}000$。
• 矩阵 $A$ 约 60%–70% 稀疏。
• 性能测试在 $M = 1{,}000,\ N = 10{,}000$ 的规模下进行。

解题思路

SpMV 的核心挑战在于不规则的内存访问。标准存储格式如 CSR（Compressed Sparse Row） 每行只存储非零元素，节省空间但导致对向量 $x$ 的间接索引访问不合并。另一种是 ELLPACK 格式，将每行填充到相同长度以改善合并访问。选择哪种格式取决于稀疏模式——如果每行非零元素数量接近，ELLPACK 更好；如果波动大，CSR 更灵活。GPU 上最优实现通常使用混合格式或 SELL-C-σ 等变体。

代码实现

CUDA

#include <cuda_runtime.h>
// CSR 格式: values[], col_indices[], row_ptr[]
__global__ void spmv_kernel(const float* values, const int* col_idx, const int* row_ptr,
    const float* x, float* y, int M) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M) {
        float sum = 0.0f;
        for (int j = row_ptr[row]; j < row_ptr[row+1]; j++)
            sum += values[j] * x[col_idx[j]];
        y[row] = sum;
    }
}
extern "C" void solve(const float* values, const int* col_idx, const int* row_ptr,
    const float* x, float* y, int M) {
    spmv_kernel<<<(M+255)/256, 256>>>(values, col_idx, row_ptr, x, y, M);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def spmv_kernel(values_ptr, col_idx_ptr, row_ptr_ptr, x_ptr, y_ptr, M: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = row < M
    start = tl.load(row_ptr_ptr + row, mask=mask)
    end = tl.load(row_ptr_ptr + row + 1, mask=mask)
    acc = tl.zeros((BLOCK,), tl.float32)
    for j in range(tl.max(end - start)):
        valid = start + j < end
        col = tl.load(col_idx_ptr + start + j, mask=mask & valid, other=0)
        val = tl.load(values_ptr + start + j, mask=mask & valid, other=0.0)
        acc += val * tl.load(x_ptr + col, mask=mask & valid, other=0.0)
    tl.store(y_ptr + row, acc, mask=mask)