Python使用CUDA

简单的numba+CUDA测试

import numpy as np
from timeit import default_timer as timer
from numba import vectorize

@vectorize(['float32(float32, float32)'], target='cuda')
def vector_add(a, b):
    return a + b


def main():
    n = 3200000

    a = np.ones(n, dtype=np.float32)
    b = np.ones(n, dtype=np.float32)
    # c = np.ones(n, dtype=np.float32)

    start = timer()
    c = vector_add(a, b)
    vector_add_time = timer() - start

    print(c[:5])
    print(c[-5:])
    print(vector_add_time)


if __name__ == '__main__':
    main()

报错：

numba.uda.cudadrv.error.NvvmSupportError: libNVVM cannot be found. Do `conda install cudatoolkit`

找了网上添加环境变量的方法无果。

之后手动添加cudatoolkit链接：

import os

os.environ['NUMBAPRO_NVVM'] = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0\nvvm\bin\nvvm64_32_0.dll'

os.environ['NUMBAPRO_LIBDEVICE'] = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0\nvvm\libdevice'

这里有个tips：

要将C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0\nvvm\libdevice中的libdevice..10.bc重命名为libdevice.compute_52.10.bc。(其中compute_52）为显卡的算力大小。不然仍然会报错。