Numba自定义CUDA设备函数能否调用numba.cuda.random设备函数

核心错误原因 #

• 设备函数的签名声明错误：你为device设备函数指定的参数类型为float32[:,:]，但实际传入的是create_xoroshiro128p_states返回的随机数状态数组，二者类型完全不匹配，导致Numba无法正确推导随机函数的输入类型，触发类型报错。
• numba.cuda.random下的随机生成函数完全支持在CUDA设备函数中调用，不存在调用权限问题。
• 额外注意：你写的核函数启动参数顺序错误，CUDA核函数启动语法为内核[块数(BPG), 每块线程数(TPB)]，你当前写反了参数顺序，即使类型错误修复后也会出现运行异常。

修复后的可运行代码 #

import numba
from numba import cuda 
from numba.cuda.random import create_xoroshiro128p_states
from numba.cuda.random import xoroshiro128p_uniform_float64

# 去掉错误的类型签名，让Numba自动推导参数类型即可
@cuda.jit(device=True) 
def device(rng_states):  
    thread_id = cuda.grid(1) 
    probability = xoroshiro128p_uniform_float64(rng_states, thread_id) 
    # 此处可补充你的业务逻辑，比如返回概率值或用概率做后续计算
     
@cuda.jit() 
def kernel(rng_states): 
    device(rng_states) 

BPG = 10 
TPB = 10 
rng_states = create_xoroshiro128p_states(BPG * TPB, seed=42069) 
# 修正核函数启动参数顺序
kernel[BPG, TPB](rng_states)

补充说明 #

如果你一定要手动写设备函数的类型签名，需要使用Numba内置的随机状态类型numba.cuda.random.xoroshiro128p_state_type[:]作为参数类型，不过更推荐直接省略类型签名让Numba自动推导，减少手写错误的概率。

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