使用带有numba和CUDA的Python解决线性系统

Question

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()

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

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.

Answer 1

Am I correct in assuming this?

No

Are there other data types that will work?

No

The problem here is that you cannot use numpy.linalg in code which is targeted to run on the numba GPU backend.