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

Question

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

我当前的非 GPU“优化”是将 A 表示为 scipy.sparse.csc_matrix 对象，然后简单地计算 A.dot(x)，但我希望在具有几个 NVIDIA 的虚拟机上加快速度连接 GPU，如果可能的话仅使用 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，因此这是一个额外的好处