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

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

0 голосов
/ 15 октября 2018

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, о котором я только мог подумать.

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

1 Ответ

0 голосов
/ 15 октября 2018

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

  1. Не украшать мои функции изначально cuda.jit, так что они все еще обладают атрибутом __name__
  2. Получение атрибута __name__
  3. Теперь применяем cuda.jit к моим функциям, напрямую вызывая декоратор
  4. Создание питона для функции композиции в строке и передавая его 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]
...