3

我正在使用 Numba（版本 0.37.0）来优化 GPU 代码。我想使用组合矢量化函数（使用 Numba 的 @vectorize 装饰器）。

进口和数据：

import numpy as np
from math import sqrt
from numba import vectorize, guvectorize

angles = np.random.uniform(-np.pi, np.pi, 10)
coords = np.stack([np.cos(angles), np.sin(angles)], axis=1)

这按预期工作：

@guvectorize(['(float32[:], float32[:])'], '(i)->()', target='cuda')
def l2_norm(vec, out):
    acc = 0.0
    for value in vec:
        acc += value**2
    out[0] = sqrt(acc)

l2_norm(coords)

输出：

array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)

但我想通过调用另一个矢量化函数来避免在“l2_norm”中使用“for”。

我试过这个：

@vectorize(["float32(float32)"], target="cuda")
def power(value):
    return value**2

@guvectorize(['(float32[:], float32[:])'], '(i)->()', target='cuda')
def l2_norm_power(vec, out):
    acc = 0.0
    acc = power(vec)
    acc = acc.sum()
    out[0] = sqrt(acc)

l2_norm_power(coords)

但引发 TypingError：

TypingError: Failed at nopython (nopython frontend)
Untyped global name 'power': cannot determine Numba type of <class  
'numba.cuda.dispatcher.CUDAUFuncDispatcher'>

关于如何执行这种组合的任何想法？

关于使用 Numba 优化 l2_norm 的其他方法有什么建议吗？

4

1 回答 1

3

我认为您只能device=True从其他 cuda 函数调用函数：

3.13.2。示例：调用设备函数

所有 CUDA ufunc 内核都能够调用其他 CUDA 设备函数：

 from numba import vectorize, cuda
 # define a device function
 @cuda.jit('float32(float32, float32, float32)', device=True, inline=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)

请注意，您可以cuda.jit使用以下方式调用函数device：

@cuda.jit(device=True)
def sum_of_squares(arr):
    acc = 0
    for item in arr:
        acc += item ** 2
    return acc

@nb.guvectorize(['(float32[:], float32[:])'], '(i)->()', target='cuda')
def l2_norm_power(vec, out):
    acc = sum_of_squares(vec)
    out[0] = sqrt(acc)

l2_norm_power(coords)

但这可能会破坏拆分它的目的。

因为numba.vectorize不支持device这些功能是不可能的。但这是一件好事，因为vectorize分配一个数组来放入值，这是一个不必要的中间数组，并且在 GPU 上分配数组也非常低效（并且在 numba 中被禁止）：

3.5.5。Numpy 支持

由于 CUDA 编程模型，内核内部的动态内存分配效率低下，通常不需要。Numba 不允许任何内存分配功能。这会禁用大量 NumPy API。为了获得最佳性能，用户应该编写代码，使每个线程一次处理一个元素。

鉴于所有这些，我将简单地使用您的原始方法：

@guvectorize(['(float32[:], float32[:])'], '(i)->()', target='cuda')
def l2_norm(vec, out):
    acc = 0.0
    for value in vec:
        acc += value**2
    out[0] = sqrt(acc)

l2_norm(coords)
于 2019-09-27T18:25:00.540 回答