使用 Python 与 numba 和 CUDA 求解线性系统

Mik*_*ike 5 python cuda numpy

I am trying to solve a linear system using numba with GPU processing using CUDA.

I have installed all the relevant packages and tested it so it seems that my GPU and CUDA etc is set up properly.

My code is:

import numpy as np
import time

from numba import vectorize, cuda


@vectorize(['float64(float64, float64)'], target='cuda')
def solver(A, b):
    return np.linalg.solve(A, b)


def main():

    A = np.random.rand(100, 100).astype(np.float64)
    b = np.random.rand(100, 1).astype(np.float64)

    start = time.time()
    C = solver(A, b)
    vector_add_time = time.time() - start

    print("Took " + str(vector_add_time) + " seconds to solve")


if __name__ == '__main__':
    main()
Run Code Online (Sandbox Code Playgroud)

Commenting the @vectorize... line, the code runs fine. However, when I try to do it with numba and cuda, I get a long list of errors, where I think he most relevant one is:

raise TypingError(msg)
numba.errors.TypingError: Failed at nopython (nopython frontend)
np.linalg.solve() only supported for array types
Run Code Online (Sandbox Code Playgroud)

I assume the problem is that numpy.linalg.solve does not accept the data types required by cuda.

Am I correct in assuming this? Are there other data types that will work?

In this example problem, the same data type is passed to the function, so I think the problem lies with numpy.linalg.

tal*_*ies 2

我的假设正确吗?

还有其他可以使用的数据类型吗?

这里的问题是您无法numpy.linalg在旨在在 numba GPU 后端运行的代码中使用。