Numba CUDA使用vectorize/guvectorize报错及返回0如何解决

常见报错原因说明 #

1. @vectorize触发float32无shape属性报错：@vectorize是元素级操作装饰器，运行时会将输入数组的单个元素传入修饰函数，因此函数内拿到的是float32标量而非数组，自然不存在shape属性。
2. @guvectorize输出始终为0：@guvectorize要求输出参数为预先分配的内存空间，若在函数内直接对输出参数做out = xxx赋值，只会修改局部变量指针，不会写入实际输出内存，因此结果始终为初始值0。
3. 1维数组声明触发索引越界：输入为2维图像，声明@guvectorize签名时若使用1维维度标记，维度不匹配会直接触发索引越界错误。

正确实现方案 #

1. 基于@guvectorize的CUDA实现 #

该方案适合快速开发，无需手动管理CUDA线程调度：

import numpy as np
from numba import guvectorize, float32, int64, cuda

# 装饰器声明：输入为两个相同尺寸的2维float32数组，输出为1维int64数组（用长度1数组避免标量赋值失效问题）
@guvectorize(
    [(float32[:,:], float32[:,:], int64[:])],
    "(n,m),(n,m)->()",
    target='cuda'
)
def count_high_pixel_diff(img1, img2, res):
    cnt1 = 0
    cnt2 = 0
    for i in range(img1.shape[0]):
        for j in range(img1.shape[1]):
            if img1[i,j] > 200:
                cnt1 += 1
            if img2[i,j] > 200:
                cnt2 += 1
    res[0] = cnt1 - cnt2 # 必须修改输出数组的元素值，不可直接对res赋值

# 调用示例
# 预先将数据加载到GPU内存，减少传输开销
img1 = cuda.to_device(np.random.uniform(0,255, size=(1024,1024)).astype(np.float32))
img2 = cuda.to_device(np.random.uniform(0,255, size=(1024,1024)).astype(np.float32))
diff = count_high_pixel_diff(img1, img2)
# 结果转回CPU输出
print(diff.copy_to_host())

2. 结合cuda.jit的手动核函数实现 #

该方案可控性更高，性能上限更高，适合对性能有要求的场景：

import numpy as np
from numba import cuda

# 线程块大小建议设为32的倍数，可根据GPU架构调整
THREADS_PER_BLOCK = 256

@cuda.jit
def count_high_pixel_diff_cuda(img1, img2, partial_cnt):
    # 共享内存存储每个线程块的局部统计结果
    shared_cnt = cuda.shared.array(THREADS_PER_BLOCK, dtype=np.int64)
    tid = cuda.threadIdx.x
    gid = cuda.grid(1)
    total_pixels = img1.shape[0] * img1.shape[1]
    
    # 初始化共享内存
    shared_cnt[tid] = 0
    cuda.syncthreads()
    
    # 网格步长遍历所有像素
    while gid < total_pixels:
        i = gid // img1.shape[1]
        j = gid % img1.shape[1]
        delta = 0
        if img1[i,j] > 200:
            delta += 1
        if img2[i,j] > 200:
            delta -= 1
        shared_cnt[tid] += delta
        gid += cuda.gridsize(1)
    
    cuda.syncthreads()
    
    # 块内归约计算
    s = THREADS_PER_BLOCK // 2
    while s > 0:
        if tid < s:
            shared_cnt[tid] += shared_cnt[tid + s]
        cuda.syncthreads()
        s //= 2
    
    # 每个块的最终结果写入全局内存
    if tid == 0:
        partial_cnt[cuda.blockIdx.x] = shared_cnt[0]

# 调用示例
img1 = cuda.to_device(np.random.uniform(0,255, size=(1024,1024)).astype(np.float32))
img2 = cuda.to_device(np.random.uniform(0,255, size=(1024,1024)).astype(np.float32))
# 计算网格尺寸
blocks_per_grid = (img1.shape[0] * img1.shape[1] + THREADS_PER_BLOCK - 1) // THREADS_PER_BLOCK
partial_cnt = cuda.to_device(np.zeros(blocks_per_grid, dtype=np.int64))
# 启动核函数
count_high_pixel_diff_cuda[blocks_per_grid, THREADS_PER_BLOCK](img1, img2, partial_cnt)
# 最终归约得到结果
total_diff = partial_cnt.copy_to_host().sum()
print(total_diff)

GPU适配优化建议 #

• 统计类操作属于典型的归约场景，优先使用块内共享内存做局部归约，可大幅降低全局内存访问开销
• 提前将输入数据转为GPU原生支持的类型（如float32、uint8），避免核函数内隐式类型转换带来的额外开销
• 尽量减少CPU与GPU之间的内存拷贝次数，所有输入输出数据批量提前转到CUDA数组，不要依赖函数隐式传输
• 涉及单个值输出的场景，优先用长度为1的数组存储结果，避免标量参数赋值不生效的问题
• 线程块大小建议设置为32的倍数，常用值为128、256、512，可根据自身GPU架构测试最优参数

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