如何在 object 中从 CUDA kernel 在 python 中调用设备 function

Question

我正在编写非常具体的神经网络，我有许多不同的激活函数类，每个类都有 function 用于正常 python 和一个作为设备 function。问题是从 CUDA kernel 内部调用该方法。

@cuda.jit(device=True)
def activation_fn(z):
    return max(0, z)

@cuda.jit
def backprop_kernel(arr):
    arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])

def backprop_GPU(x, y):
    arr = np.array([-3, -2, -1, 0, 1, 2, 3])
    print(arr)
    backprop_kernel[1, 7](arr)
    print(arr)

backprop_GPU(None, None)

这工作得很好，但我想让下面的代码工作。

class Activation:
    
    @cuda.jit(device=True)
    def fn(z):
        return max(0, z)

class Network:

    def __init__(self):
        self.activation_fn = Activation()
    
    @cuda.jit
    def kernel(arr):
        arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])

    def backprop(self, x, y):
        arr = np.array([-3, -2, -1, 0, 1, 2, 3])
        self.kernel[1, 7](arr)

net = Network()
net.backprop(None, None)

如何使“activation_fn”可从 kernel 访问？

Answer 1

@cuda.jit必须与函数一起使用，而不是成员，因此您需要在方法内部定义修饰函数，并在定义 kernel 时捕获激活 function：

from numba import cuda
import numpy as np


class Activation:
    def __init__(self):
        @cuda.jit(device=True)
        def fn(z):
            return max(0, z)

        self.fn = fn


class Network:
    def __init__(self):
        self.activation = Activation()

        activation_fn = self.activation.fn

        @cuda.jit
        def kernel(arr):
            arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])

        self.kernel = kernel

    def backprop(self, x, y):
        arr = np.array([-3, -2, -1, 0, 1, 2, 3])
        self.kernel[1, 7](arr)
        print(arr)


net = Network()
net.backprop(None, None)

印刷：

$ python repro.py 
[0 0 0 0 1 2 3]

（请注意，我省略了此处出现的性能警告，因为它们与手头的问题正交）