CuPy中编写支持cuSPARSE实例的自定义内核API及CUDA参数问询

要编写能处理cuSPARSE兼容稀疏矩阵的自定义内核，CuPy提供两类核心API：

• cupy.ElementwiseKernel：适合元素级运算，但对稀疏矩阵原生支持有限，一般需先提取非零元素为稠密数组处理，再重构稀疏矩阵；若要直接操作稀疏结构，更推荐使用RawKernel。
• cupy.RawKernel：灵活性更高，允许直接编写底层CUDA C代码，可直接访问稀疏矩阵的底层存储结构，完美适配cuSPARSE实例类型。

当编写RawKernel接收csr_matrix作为输入时，底层CUDA代码需要接收CSR格式的核心存储组件与元数据参数：

• 核心存储组件（对应csr_matrix的属性）：
  • data：非零元素数组，对应CUDA中的const <数据类型>*（如const float*）
  • indices：非零元素的列索引数组，对应CUDA中的const int*
  • indptr：行指针数组，标记每行非零元素在data/indices中的起止位置，对应CUDA中的const int*
• 辅助元数据：
  • n_rows：矩阵行数，int类型
  • n_cols：矩阵列数，int类型
  • nnz：非零元素总数，int类型

简单示例框架 #

CUDA内核代码：

extern "C" __global__ void csr_process(const float* data, const int* indices, const int* indptr, int n_rows, int nnz) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < nnz) {
        // 访问第idx个非零元素的值：data[idx]
        // 对应的列索引：indices[idx]
        // 可通过indptr推导元素所属行
    }
}

Python侧调用：

import cupy as cp
from cupyx.scipy.sparse import csr_matrix

# 构造示例CSR矩阵
csr = csr_matrix((cp.array([1,2,3]), cp.array([0,2,1]), cp.array([0,2,3])), shape=(2,3))

# 定义RawKernel
kernel = cp.RawKernel(r'''
extern "C" __global__ void csr_process(const float* data, const int* indices, const int* indptr, int n_rows, int nnz) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < nnz) {
        // 自定义处理逻辑
        printf("Element %d: value=%f, col=%d\n", idx, data[idx], indices[idx]);
    }
}
''', 'csr_process')

# 调用内核
kernel((1,), (csr.nnz,), (csr.data, csr.indices, csr.indptr, csr.shape[0], csr.nnz))

内容的提问来源于stack exchange，提问作者linkhyrule5