如何使用 CUDA 通过当前通过 scipy.sparse.csc_matrix.dot 实现的密集向量积来加速稀疏矩阵？

Question

我的最终目标是在 Python 中加速矩阵向量乘积的计算，可能是通过使用支持 CUDA 的 GPU。矩阵 A 约为 15k x 15k 且稀疏（密度 ~ 0.05），向量 x 为 15k 元素且密集，我正在计算 Ax。我必须多次执行此计算，因此使其尽可能快是理想的。

我当前的非 GPU“优化”是将 A 表示为 scipy.sparse.csc_matrix 对象，然后简单地计算 A.dot(x)，但我希望在连接了几个 NVIDIA GPU 的 VM 上加快速度，并且尽可能只使用 Python（即不手动写出详细的内核函数）。我已经使用 cudamat 库成功地加速了密集矩阵向量产品，但不适用于稀疏情况。网上有一些针对稀疏案例的建议，例如使用pycuda，或scikit-cuda，或anaconda 的加速包，但没有大量信息，因此很难知道从哪里开始。

我不需要非常详细的说明，但是如果有人以前解决过这个问题并且可以为最简单的方法提供一个“大图”路线图，或者有一个基于 GPU 的稀疏矩阵加速的想法 -矢量积将超过 scipy 的稀疏算法，这将非常有帮助。

Answer 1

另一种选择是使用CuPy包。它具有与 numpy/ scipy 相同的界面（这很好）并且（至少对我来说），事实证明它比pycuda. 新代码看起来像这样：

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

A = some_sparse_matrix #(scipy.sparse.csr_matrix)
x = some_dense_vector  #(numpy.ndarray)

A_gpu = csr_gpu(A)  #moving A to the gpu
x_gpu = cp.array(x) #moving x to the gpu

for i in range(niter):
    x_gpu = A_gpu.dot(x_gpu)
x = cp.asnumpy(x_gpu) #back to numpy object for fast indexing

更新：cupy 现在也通过其 ROCm 堆栈支持 AMD GPU，因此这是一个额外的好处

Answer 2

正如评论中指出的，NVIDIA 提供了cuSPARSE库，其中包含稀疏矩阵乘积与稠密向量的函数。

Numba 现在通过pyculib包具有 cuSparse 库的 Python 绑定。