如何在Python的CUDA核函数中调用对象内的设备函数？

我正在编写一个包含多种激活函数类的神经网络，每个类都实现了普通Python函数和JIT编译的设备函数。现在遇到的问题是如何在CUDA核函数内部调用类的方法。

以下代码可以正常运行：

from numba import cuda
import numpy as np

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

但我希望让以下代码也能正常运行：

from numba import cuda
import numpy as np

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？

由于Numba的CUDA核函数无法直接捕获类实例的上下文，编译阶段无法识别实例属性，可通过以下几种方式解决：

方法1：将设备函数作为参数传入核函数 #

Numba允许将设备函数作为参数传递给核函数，调用时直接传入激活函数的设备方法即可：

from numba import cuda
import numpy as np

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

class Network:
    def __init__(self):
        self.activation_fn = Activation.fn  # 直接引用静态设备函数
    
    @cuda.jit
    def kernel(arr, act_fn):
        idx = cuda.threadIdx.x
        arr[idx] = act_fn(arr[idx])

    def backprop(self, x, y):
        arr = np.array([-3, -2, -1, 0, 1, 2, 3])
        print("Before:", arr)
        # 调用核函数时传入激活函数
        self.kernel[1, 7](arr, self.activation_fn)
        print("After:", arr)

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

注意需将激活函数定义为静态方法，避免实例化依赖。

方法2：直接调用类的静态设备函数 #

如果激活函数不需要实例状态，可直接将其定义为类的静态方法，核函数中通过类名直接调用：

from numba import cuda
import numpy as np

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

class Network:
    @cuda.jit
    def kernel(arr):
        idx = cuda.threadIdx.x
        arr[idx] = Activation.relu(arr[idx])

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

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

这种方式适合无状态的激活函数，实现简单直接。

方法3：使用jitclass封装带状态的激活函数 #

如果激活函数需要实例参数（如LeakyReLU的alpha参数），可使用Numba的jitclass定义可编译的激活类，将实例传递到核函数中：

from numba import cuda, jitclass, float32

# 定义jitclass的属性规范
activation_spec = [
    ('alpha', float32),  # LeakyReLU的斜率参数
]

@jitclass(activation_spec)
class LeakyReLU:
    def __init__(self, alpha):
        self.alpha = alpha
    
    def fn(self, z):
        return max(self.alpha * z, z)

# 包装jitclass方法为设备函数
@cuda.jit(device=True)
def apply_activation(act_obj, z):
    return act_obj.fn(z)

class Network:
    def __init__(self):
        self.activation = LeakyReLU(0.1)
    
    @cuda.jit
    def kernel(arr, act_obj):
        idx = cuda.threadIdx.x
        arr[idx] = apply_activation(act_obj, arr[idx])

    def backprop(self, x, y):
        arr = np.array([-3, -2, -1, 0, 1, 2, 3], dtype=np.float32)
        print("Before:", arr)
        # 传递jitclass实例到核函数
        self.kernel[1, 7](arr, self.activation)
        print("After:", arr)

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

该方案支持带实例状态的激活函数，适合复杂场景。

内容的提问来源于stack exchange，提问作者Jirka Klimeš