Numba/CUDA - 调用矢量化库函数

Question

我想使用 Numba 让 GPU 运行定义了以下函数 here，它具有签名：

@numba.vectorize(_signatures)
def ppf(p, df, mu, sigma):
    # do stuff here

我不确定这是否相关，但我注意到因为这个函数是在单独的包中定义的，所以我无法在装饰器中指定 target='cuda'。

例如，考虑以下片段：

from numba import cuda, vectorize

@cuda.jit('float32(float32, float32, float32)', device=True)
def cu_device_fn(x, y, z):
    return x ** y / z

# define a ufunc that calls our device function
@vectorize(['float32(float32, float32, float32)'], target='cuda')
def cu_ufunc(x, y, z):
    return cu_device_fn(x, y, z)

直截了当，我可以直接通过cu_ufunc(10, 2, 2)调用这个函数。

我如何做同样的事情，即让 GPU 运行这个导入的 Numba 函数？

很抱歉缺少详细信息。我是 GPU 编程的完全初学者，完全迷失了。

Answer 1

I'm not sure if this is relevant, but I noticed that because this function is defined in a separate package, I can't specify in the decorator target='cuda'.

问题不在于该函数是在另一个模块中定义的，而是该函数已经被修饰了。导入后不能替换装饰器，因为函数已经定义（并被装饰实现替换）。

当你看到这样的东西时：

@numba.vectorize(_signatures)
def ppf(p, df, mu, sigma):
    # do stuff here

真正发生的事情是这样的：

def ppf(p, df, mu, sigma):
    # do stuff here

ppf = numba.vectorize(ppf, _signatures)

即原始 Python 被传递给矢量化系统，并将编译后的矢量化代码的结果包装器分配给函数。原始代码不再存在（除了解释器字节码的某个内部版本），因此无法再次修饰它。

您必须分叉代码并修改它以满足您的目的。请注意，完全可以这样做：

# define a ufunc that calls our device function
def _ppf(p, df, mu, sigma):
    # do stuff here

ppf = numba.vectorize(_ppf, _signatures)
ppf_gpu = numba.vectorize(_ppf, _signatures, target='cuda')

这为您提供了来自原始源的两个不同的 JIT 函数，用于不同的目标。

后记：这里暗示无论 # do stuff here 是什么，它都将是纯 python 代码，它只使用 Numba 可以为 CUDA 和 CPU 矢量化目标。如果不是这种情况，那么此答案对您没有任何帮助。

Numba/CUDA - 调用矢量化库函数

Numba/CUDA - Calling vectorized library function

python

cuda

gpu

numba