Codon GPU编程完全教程：用@gpu.kernel在Python里跑CUDA核函数绘制Mandelbrot集

Codon GPU编程完全教程用gpu.kernel在Python里跑CUDA核函数绘制Mandelbrot集【免费下载链接】codonA high-performance, zero-overhead, extensible Python compiler with built-in NumPy support项目地址: https://gitcode.com/gh_mirrors/co/codonCodon 是一个高性能、零开销、可扩展的 Python 编译器内置 NumPy 支持能直接编译出原生机器码。本教程带你用 Codon 的GPU 编程能力仅凭一行gpu.kernel装饰器在 Python 里编写并运行真正的CUDA 核函数kernel最终在 GPU 上并行计算并绘制经典的Mandelbrot 集曼德博集合分形图像——相比 CPU 版本可快约450 倍。一、Codon GPU 编程为什么值得学Codon 的口号是把 Python 重新想象为面向静态、提前编译的语言。单线程下相比原版 Python 通常有 10–100 倍甚至更高的加速并且原生支持多核多线程没有 GIL和GPU 编程。官方文档明确列出的核心用途包括在 Python 里写多线程或 GPU 代码加速并并行化 NumPy 代码把 Python 编译到边缘设备、嵌入式系统等硬件上跑在 GPU 方面Codon 提供原生 GPU 后端目前支持Nvidia 设备需要安装 CUDA 工具包Codon 会自动加载libcuda.somacOS 为libcuda.dylib若库不在标准位置可用环境变量CODON_CUDA指定路径生成的代码最终变成 PTX内嵌为常量数据运行时通过 CUDA Driver API 加载完整 API 文档见 docs/parallel/gpu.mdGPU 代码生成的编译器实现位于 codon/cir/llvm/gpu.cpp标准库的gpu模块入口在 stdlib/gpu.codon底层 CUDA 封装见 stdlib/internal/gpu.codon。二、环境准备3 步跑通 Codon GPU 环境安装 Codon按 docs/start/install.md 的说明安装获得codon命令安装 CUDA ToolkitNvidia GPU 编程的必备依赖验证设备Codon 运行时通过 CUDA Driver API 探测设备可用--gpu-name编译选项指定目标 GPU 架构关于编译选项Codon 提供两个关键开关详见 docs/parallel/gpu.md选项作用默认值--gpu-name name指定目标 GPU 架构 / compute capability直接传给 LLVMsm_30--gpu-features features以逗号分隔的 LLVM 风格 GPU 特性开关如ptx42ptx42 建议如果你的显卡较新如 RTX 30/40 系把--gpu-name设成对应的sm_86/sm_89能获得更好的指令集支持。三、第一个 GPU 核函数gpu.kernel 入门核函数kernel就是运行在 GPU 上的函数用gpu.kernel装饰标记Codon 会对其进行特殊编译。下面是最小的向量加法核函数import gpu gpu.kernel def hello(a, b, c): i gpu.thread.x c[i] a[i] b[i] a [i for i in range(16)] b [2*i for i in range(16)] c [0 for _ in range(16)] hello(a, b, c, grid1, block16) print(c) # [0, 3, 6, ..., 45]代码几乎就是纯 Python定义核函数、准备数据、调用时额外传grid和block参数即可。调用完成后结果自动拷回 CPU直接print就能用。内置 GPU 内建量intrinsicCodon 的gpu模块提供了与 CUDA 几乎一一对应的内建量Codon含义CUDA 等价物gpu.thread.x当前线程在块内的 x 坐标threadId.xgpu.block.x当前块在网格中的 x 坐标blockIdx.xgpu.block.dim.x块的 x 维大小blockDim.xgpu.grid.dim.x网格的 x 维大小gridDim.xy、z坐标同理*.dim对象都是gpu.Dim3实例。四、理解 GPU 执行模型grid 与 block 怎么设GPU 执行时程序运行在一个网格grid上网格由 (X × Y × Z) 个块block组成每个块又包含 (x × y × z) 个线程thread。调用核函数时grid和block参数支持以下四种写法单个整数x→ 维度(x, 1, 1)两个整数元组(x, y)→ 维度(x, y, 1)三个整数元组(x, y, z)→ 完整三维gpu.Dim3实例如Dim3(x, y, z)例如hello(a, b, c, grid1, block16)表示1 个块、16 个线程正好处理 16 个元素。注意具体设备的 grid / block 尺寸上限会受硬件限制。 经验法则处理 N 个独立元素时常取block256 或 1024硬件友好的对齐值grid (N block - 1) // block向上取整覆盖全部数据。五、实战GPU 核函数绘制 Mandelbrot 集Mandelbrot 集的分形图像是 GPU 并行的经典测试题每个像素独立迭代z z² c天然适合一个线程算一个像素。下面是官方文档给出的完整示例from python import matplotlib.pyplot as plt import numpy as np import gpu MAX 1000 # 最大迭代次数 N 4096 # 图像宽和高 pixels np.empty((N, N), dtypeint) def scale(x, a, b): return a (x/N)*(b - a) gpu.kernel def mandelbrot(pixels): idx (gpu.block.x * gpu.block.dim.x) gpu.thread.x i, j divmod(idx, N) c complex(scale(j, -2.00, 0.47), scale(i, -1.12, 1.12)) z 0j iteration 0 while abs(z) 2 and iteration MAX: z z**2 c iteration 1 pixels[i, j] int(255 * iteration/MAX) mandelbrot(pixels, grid(N*N)//1024, block1024) plt.imshow(pixels) plt.show()逐行拆解几个关键点全局线程编号idx gpu.block.x * gpu.block.dim.x gpu.thread.x把块坐标和线程坐标拼成唯一的像素索引这是 CUDA 编程的标准套路在 Codon 里写法和 C 一模一样divmod解包i, j divmod(idx, N)把一维索引还原成二维行列坐标复数运算直接可用complex类型在核函数内开箱即用迭代循环z z**2 c无需任何改写NumPy 无缝支持pixels是 NumPy 数组直接以pixels[i, j]二维下标读写数据会自动拷贝到 GPU、算完再拷回绘图结果回到 CPU 后直接from python import matplotlib调用 Matplotlib 可视化官方实测这段 GPU 版本比等价的 CPU 版本快约 450 倍。仓库里的性能基准 bench/codon/mandelbrot.codon 使用了同样的 Mandelbrot 算法见 bench/README.md你甚至可以codon run bench/codon/mandelbrot.codon亲手计时。六、更省事的方式par(gpuTrue) 自动并行化如果你不想手写核函数Codon 的par语法可以把现有循环直接搬到 GPU 上一行装饰器搞定。同一幅 Mandelbrot 图可以这样生成import numpy as np MAX 1000 N 4096 pixels np.empty((N, N), dtypeint) def scale(x, a, b): return a (x/N)*(b - a) par(gpuTrue, collapse2) for i in range(N): for j in range(N): c complex(scale(j, -2.00, 0.47), scale(i, -1.12, 1.12)) z 0j iteration 0 while abs(z) 2 and iteration MAX: z z**2 c iteration 1 pixels[i, j] int(255 * iteration/MAX)要点gpuTrue把整个循环并行化到 GPU无需手动计算 grid/blockcollapse2用于循环嵌套场景把两层循环的迭代空间压平成一个整体注意gpuTrue模式不允许共享变量循环体内给循环外变量赋值和规约操作对于单维数组运算甚至可以直接a np.arange(16) b np.arange(16) * 2 c np.empty(16, dtypeint) par(gpuTrue) for i in range(16): c[i] a[i] b[i]NumPy 数组和运算在 GPU 代码里无缝工作这让加速现有 NumPy 循环变成零成本操作。七、进阶技巧数学函数、原始指针与对象转换数学函数自动替换为 GPU 优化版本math模块的所有函数在核函数中都能用并被自动替换为 GPU 优化版本底层调用 libdeviceimport math import gpu gpu.kernel def hello(x): i gpu.thread.x x[i] math.sqrt(x[i]) # 实际使用 libdevice 的 __nv_sqrt x [float(i) for i in range(10)] hello(x, grid1, block10)NumPy 的数学函数同理。libdevice 默认路径为/usr/local/cuda/nvvm/libdevice/libdevice.10.bc也可用-libdevice编译选项指定。gpu.raw原始指针模式默认情况下传入核函数的对象会完整转换为 GPU 副本保留 Python/Codon 的全部语义。如果想更接近 C/CUDA 的裸指针风格可以用gpu.raw# 以三个 int 指针参数方式调用核函数 hello(gpu.raw(a), gpu.raw(b), gpu.raw(c), grid1, block16)好处是省掉一次指针间接寻址代价是核函数参数只剩Ptr的有限 API主要是索引/赋值不再是完整的列表 API。对象如何往返 GPUCodon 用两个魔法方法在 CPU 与 GPU 之间搬运对象__to_gpu__(self)分配 GPU 内存并把对象拷到设备__from_gpu__(self, gpu_object)把 GPU 内存拷回 CPU 对象内置类型全都自带这两个方法用户自定义类也会自动生成所以大多数对象可以无缝往返 GPU。若你的类用了裸指针等底层构造需要参考 stdlib/internal/gpu.codon 中的实现自行定义。八、避坑指南核函数里的限制与排错核函数支持 Codon 的绝大部分特性但有几条重要限制也是最常见的坑限制项说明❌ 异常处理核函数内不能抛出/捕获异常raise会被标记为不可达并优化掉❌ I/O 操作不能在核函数里打开文件等❌ 部分模块re依赖外部正则库、os等模块不可用排错建议CUDA 报错核函数中触发 CUDA 错误如误用异常会被打印出来通常意味着核函数代码里有非法操作先对照上表自查查看生成的 PTX用-ptx file编译选项把生成的 PTX 存到文件里检查便于和 CUDA 行为对拍指定架构不匹配出现指令集相关报错时检查--gpu-name是否与实际显卡的 compute capability 对应更多内核行为示例可参考测试用例 test/transform/kernels.codon其中覆盖了原始指针、对象转换、用户类含dataclass(gpuTrue)等场景。九、学习路线与参考资料 建议按以下顺序上手读 docs/parallel/gpu.md 掌握全部 GPU API 语义跑通第三节的最小向量加法核函数复现第五节 Mandelbrot 示例用-ptx观察生成的 PTX尝试把gpu.kernel版本改写成par(gpuTrue, collapse2)版本对比跑基准 bench/codon/mandelbrot.codon 实测你的 GPU 加速比相关源码与文档索引GPU 编程文档docs/parallel/gpu.md多线程对照 GPU 理解 CPU 并行docs/parallel/multithreading.md标准库gpu模块stdlib/gpu.codonCUDA Driver API 封装stdlib/internal/gpu.codonGPU 代码生成实现codon/cir/llvm/gpu.cpp / codon/cir/llvm/gpu.hGPU 核函数测试test/transform/kernels.codon性能基准bench/README.md从一个装饰器到450 倍加速的分形图像Codon 让 GPU 编程从 C/CUDA 的陡峭曲线变成了几乎和写普通 Python 循环一样自然的体验。【免费下载链接】codonA high-performance, zero-overhead, extensible Python compiler with built-in NumPy support项目地址: https://gitcode.com/gh_mirrors/co/codon创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考