如何在Numba JIT GPU后端中使用列表作为函数输入？

Numba GPU后端（如numba.cuda.jit）不支持直接传入原生Python列表，因为GPU需要连续的内存布局和同质类型的数据。必须将列表转换为Numba/GPU兼容的数组结构，以下分两种常见场景给出实现方案：

场景1：列表元素是同质数值组合 #

如果objects中的每个元素是固定结构的数值集合（比如包含位置、半径的四元组），可以用NumPy结构化数组来定义类型：

步骤1：定义数据类型 #

import numba
from numba import cuda
import numpy as np

# 定义每个object的结构化dtype（字段名+类型）
np_obj_dtype = np.dtype([
    ('x', np.float32),
    ('y', np.float32),
    ('z', np.float32),
    ('radius', np.float32)
])

步骤2：转换Python列表为结构化数组 #

# 原Python列表示例
objects_py = [
    (1.0, 2.0, 3.0, 0.5),
    (4.0, 5.0, 6.0, 1.0),
    # 更多元素...
]

# 转换为GPU兼容的结构化数组
objects_arr = np.array(objects_py, dtype=np_obj_dtype)

步骤3：编写GPU JIT函数 #

@cuda.jit
def ray(objects, output):
    # 获取当前线程索引
    idx = cuda.grid(1)
    if idx >= output.size:
        return
    
    # 访问数组中的元素字段
    obj = objects[idx]
    x = obj['x']
    y = obj['y']
    radius = obj['radius']
    
    # 此处编写你的射线检测逻辑
    output[idx] = x + y + radius  # 示例计算

步骤4：调用GPU函数 #

# 准备输出数组
output = np.zeros(len(objects_arr), dtype=np.float32)
# 配置线程块与网格大小
threads_per_block = 256
blocks_per_grid = (len(objects_arr) + threads_per_block - 1) // threads_per_block
# 执行GPU函数
ray[blocks_per_grid, threads_per_block](objects_arr, output)
# 取回计算结果
result = output.copy()

场景2：列表元素是自定义类 #

如果objects是自定义Python类的实例列表，需要先将类转换为Numba jitclass，再将实例属性拆分为单独的数组（GPU对拆分后的数组访问效率更高）：

步骤1：定义jitclass #

@numba.jitclass([
    ('x', numba.float32),
    ('y', numba.float32),
    ('z', numba.float32),
    ('radius', numba.float32),
    ('color', numba.uint8[:])  # 支持数组类型属性
])
class Object:
    def __init__(self, x, y, z, radius, color):
        self.x = x
        self.y = y
        self.z = z
        self.radius = radius
        self.color = color

步骤2：拆分实例属性为单独数组 #

# 创建jitclass实例列表
objects_py = [
    Object(1.0, 2.0, 3.0, 0.5, np.array([255,0,0], dtype=np.uint8)),
    Object(4.0, 5.0, 6.0, 1.0, np.array([0,255,0], dtype=np.uint8))
]

# 将属性拆分为独立数组（GPU访问更高效）
x_arr = np.array([obj.x for obj in objects_py], dtype=np.float32)
y_arr = np.array([obj.y for obj in objects_py], dtype=np.float32)
radius_arr = np.array([obj.radius for obj in objects_py], dtype=np.float32)
color_arr = np.array([obj.color for obj in objects_py], dtype=np.uint8)

步骤3：编写GPU JIT函数 #

@cuda.jit
def ray(x_arr, y_arr, radius_arr, color_arr, output):
    idx = cuda.grid(1)
    if idx >= output.size:
        return
    
    x = x_arr[idx]
    y = y_arr[idx]
    radius = radius_arr[idx]
    red = color_arr[idx, 0]
    
    # 编写射线逻辑
    output[idx] = (x + y) * radius + red

核心注意事项 #

• 必须使用连续内存的同质数组替代原生Python列表，GPU不支持动态类型的列表结构
• 优先选择拆分属性为单独数组的方式，比结构化数组的内存访问效率更高
• 禁止在GPU函数中使用numba.typed.List，GPU对动态列表的支持有限，会导致性能骤降

内容的提问来源于stack exchange，提问作者TheDumbCoder