我正在使用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的其他方法的任何建议?
我认为您只能从其他cuda函数调用device=True
函数:
3.13.2. Example: Calling Device Functions
所有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中被禁止):
由于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)