Numba CUDA弃用错误：如何更新曼德博集合GPU代码？

问题根源 #

你的代码出现LoweringError主要有以下几个原因：

• 原代码先定义CPU版本的mandel函数（用@jit装饰），再尝试转成CUDA设备函数，Numba不支持这种跨编译路径的直接转换，CPU JIT函数无法被CUDA设备函数直接复用。
• 缺少必要的依赖库导入，导致后续运行环节报错。
• 返回值类型与图像数组类型不匹配，可能引发类型转换错误。

修正后的完整代码 #

# 导入必要依赖库
import numpy as np
from numba import cuda
from timeit import default_timer as timer
import matplotlib.pyplot as plt

# 直接定义CUDA设备函数
@cuda.jit(device=True)
def mandel(x, y, max_iters):
    """
    判断复数是否属于曼德博集合（CUDA设备端实现）
    """
    c = complex(x, y)
    z = 0.0j
    for i in range(max_iters):
        z = z * z + c
        if (z.real * z.real + z.imag * z.imag) >= 4:
            return i
    return max_iters

# 定义CUDA核函数
@cuda.jit
def mandel_kernel(min_x, max_x, min_y, max_y, image, iters):
    # 获取当前线程的二维索引
    x, y = cuda.grid(2)
    width = image.shape[1]
    height = image.shape[0]
    
    # 确保线程操作的像素在图像范围内
    if x < width and y < height:
        pixel_size_x = (max_x - min_x) / width
        pixel_size_y = (max_y - min_y) / height
        real = min_x + x * pixel_size_x
        imag = min_y + y * pixel_size_y
        # 限制返回值在uint8范围内，避免溢出
        image[y, x] = min(mandel(real, imag, iters), 255)

# 主执行逻辑
if __name__ == "__main__":
    # 创建空图像数组，指定uint8类型
    gimage = np.zeros((1024, 1536), dtype=np.uint8)
    # 定义线程块大小
    blockdim = (32, 8)
    # 根据图像尺寸自动计算网格大小，确保覆盖所有像素
    griddim = (
        (gimage.shape[1] + blockdim[0] - 1) // blockdim[0],
        (gimage.shape[0] + blockdim[1] - 1) // blockdim[1]
    )
    
    start = timer()
    # 将主机数据拷贝到CUDA设备
    d_image = cuda.to_device(gimage)
    # 启动CUDA核函数
    mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20)
    # 将设备计算结果拷贝回主机
    d_image.to_host()
    dt = timer() - start
    
    print(f"Mandelbrot created on GPU in {dt:.6f} s")
    
    # 渲染并显示图像
    plt.imshow(gimage, cmap='viridis')
    plt.show()

关键修正点 #

• 设备函数定义：直接用@cuda.jit(device=True)装饰mandel函数，跳过CPU JIT编译环节，确保函数能被CUDA核函数正确调用。
• 核函数优化：采用线程索引直接映射像素的方式，替代原嵌套循环，更符合CUDA并行编程模型，提升执行效率。
• 类型安全：通过min(..., 255)限制返回值在uint8的取值范围内，避免类型溢出错误。
• 自动网格计算：根据图像尺寸和线程块大小自动计算网格维度，确保所有像素都被线程覆盖。

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