Комбинированные векторизованные функции в Numba

Question

Я использую Numba (версия 0.37.0) для оптимизации кода для графического процессора.Я хотел бы использовать комбинированные векторизованные функции (используя @vectorize декоратор Numba).

Импорт и данные:

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)

Но я бы хотел не использовать это "for" внутри "l2_norm", вызывая другую векторизованную функцию.

Я пробовал это:

@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'>

Есть идеи о том, как выполнить эту комбинацию?

Есть предложения о другом способе оптимизации l2_norm с помощью Numba?

Answer 1

Я думаю, вы можете вызывать device=True функции только из других функций cuda:

3.13.2.Пример: вызов функций устройства

Все ядра ufunc CUDA имеют возможность вызывать другие функции устройства 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 запрещает любые функции выделения памяти.Это отключает большое количество API-интерфейсов NumPy.Для лучшей производительности пользователи должны писать код так, чтобы каждый поток имел дело с одним элементом за раз.

---

Учитывая все это, я бы просто использовал ваш оригинальный подход:

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