Программные вложенные вызовы функций numba.cuda

Question

Программные вложенные вызовы функций numba.cuda

1

Numba & CUDA noob здесь. Я бы хотел, чтобы одна функция numba.cuda программным numba.cuda другую из устройства, не передавая данные на хост. Например, учитывая настройку

from numba import cuda

@cuda.jit('int32(int32)', device=True)
def a(x):
    return x+1

@cuda.jit('int32(int32)', device=True)
def b(x):
    return 2*x

Я хотел бы иметь возможность определять составную функцию ядра, такую как

@cuda.jit('void(int32, __device__, int32)')
def b_comp(x, inner, result):
    y = inner(x)
    result = b(y)

и успешно получить

b_comp(1, a, result)
assert result == 4

В идеале я хотел бы, чтобы b_comp принимал переменные аргументы функции после компиляции [например, после вышеупомянутого вызова, чтобы все еще принимать b_comp(1, b, result) ], но решение, в котором аргументы функции фиксируются во время компиляции, по-прежнему будет работать для меня.

Из того, что я прочитал, кажется, что CUDA поддерживает передачу указателей функций. Этот пост предполагает, что numba.cuda не имеет такой поддержки, но пост не убедителен, и он также год. На странице поддерживаемого Python в numba.cuda не упоминается поддержка указателя на функцию. Но он ссылается на поддерживаемый Python на странице numba, что дает понять, что numba.jit() поддерживает функции как аргументы, хотя они фиксируются во время компиляции. Если numba.cuda.jit() делает то же самое, как я сказал выше, это сработает. В этом случае при указании подписи для comp, как я должен указывать тип переменной? Или я могу использовать numba.cuda.autojit()?

Если numba не поддерживает такой прямой подход, метапрограммирует разумный вариант? Например, когда я знаю inner функцию, мой скрипт может создать новый скрипт, содержащий функцию python, которая объединяет эти конкретные функции, а затем применить numba.cuda.jit(), а затем импортировать результат. Это кажется запутанным, но это единственный другой параметр numba -based, о котором я мог думать.

Если numba не будет делать трюк вообще или, по крайней мере, не без серьезной кладке, я был бы доволен ответом, который дал несколько подробностей, а также повторить "переход на PyCuda".

Twiffy 15 окт. 2018, в 00:43

Источник

Теги:

python

cuda

numba

1 ответ

Ещё вопросы

Twiffy · Answer 1 · 2018-10-15T00-20-00.000Z

Вот что сработало для меня:

Не украшая мои функции cuda.jit изначально, так что они все еще обладают атрибутом __name__
Получение атрибута __name__
Теперь применив cuda.jit к моим функциям, напрямую вызвав декоратор
Создание python для функции композиции в строке и передача ее в exec

Точный код:

from numba import cuda
import numpy as np


def a(x):
    return x+1

def b(x):
    return 2*x


# Here, pretend we've been passed the inner function and the outer function as arguments
inner_fun = a
outer_fun = b

# And pretend we have noooooo idea what functions these guys actually point to
inner_name = inner_fun.__name__
outer_name = outer_fun.__name__

# Now manually apply the decorator
a = cuda.jit('int32(int32)', device=True)(a)
b = cuda.jit('int32(int32)', device=True)(b)

# Now construct the definition string for the composition function, and exec it.
exec_string = '@cuda.jit(\'void(int32, int32[:])\')\n' \
              'def custom_comp(x, out_array):\n' \
              '    out_array[0]=' + outer_name + '(' + inner_name + '(x))\n'

exec(exec_string)

out_array = np.array([-1])
custom_comp(1, out_array)
print(out_array)

Как и ожидалось, выход

[4]