编译器高性能计算【免费下载链接】numbaNumPy aware dynamic Python compiler using LLVM项目地址https://gitcode.com/gh_mirrors/nu/numba点击查看免费下载Numba 的 CUDA 后端允许 Python 编写的 kernel 直接调用用其他语言编写的设备函数例如 CUDA C/C、PTX 汇编以及预编译的二进制对象cubins、fat binaries 等。本文以官方文档 docs/source/cuda/cuda_ffi.rst 为核心结合 numba/cuda/decorators.py、numba/cuda/cudadrv/driver.py 等源码系统讲解设备函数 ABI 约定、Python 侧声明、指针传递、链接与调用流程并给出可直接运行的完整示例。读完本文你将能够在 Numba CUDA kernel 中无缝复用已有的 CUDA C/C 设备函数库。一、总体架构Python Kernel 调用外部设备函数的三个组成部分从文档可知一次 Python kernel 对“外部设备函数”的调用由三部分组成外部设备函数的实现用其他语言如 CUDA C编写最终必须能编译为 PTX 或直接提供二进制对象cubin、fatbin 等。设备函数在 Python 中的声明通过cuda.declare_device()为外部函数声明签名生成可被 kernel 调用的描述符。与外部函数链接并调用它的 kernel通过cuda.jit(link[...])将外部函数链接进 kernel然后在 kernel 内部像调用普通函数一样调用它。支持矩阵如下外部代码形式是否直接支持说明CUDA C/C 源码.cu支持由 NVRTC 编译为 PTX 后链接PTX 源码.ptx支持直接交给 CUDA Linker 链接二进制对象cubin、fatbin 等支持直接交给 CUDA Linker 链接其他语言源码需先编译为 PTX编译后的 PTX 再按上述方式链接二、设备函数 ABIC 侧必须遵循的原型规范Numba 为调用设备函数定义了一套严格的 ABIApplication Binary InterfaceC/C 侧的原型如下extern C __device__ int function( T* return_value, ... );原型各组成部分的含义见 cuda_ffi.rst 的 Device function ABI 一节extern C防止 C 名称修饰name mangling使 Python 侧声明外部函数名时无需处理修饰后的名称。可以去掉但去掉后 Python 声明中必须使用修饰后的mangled名称。__device__必须保留用于将该函数定义为设备函数device function。返回值类型固定为int该返回值用于向调用方Numba 运行时传递Python 异常是否发生的信号。外部函数中不会发生 Python 异常因此 callee被调用方应始终返回 0。第一个参数是指向返回值的指针T* return_value由调用方分配位于 local 地址空间见下文注释并传入被调用函数。如果函数有返回值callee 应将结果写入*return_value。其余参数类型和顺序必须与 Python kernel 调用该函数时传入的参数一一对应。注意文档脚注return_value指向的数据位于local 地址空间。对 local 地址空间中的数据做原子操作atomics等操作是不适用的因此任何对返回值执行的操作都必须确保对 local 地址空间数据合法。其他语言编写的函数必须编译成符合该原型规范的 PTX才能被 Numba 正确链接和调用。一个具体的原型示例一个接受两个 float、返回一个 float 的函数的原型如下extern C __device__ int mul_f32_f32( float* return_value, float x, float y );该示例的完整实现可参见仓库中的 numba/cuda/tests/doc_examples/ffi/functions.cuextern C __device__ int mul_f32_f32( float* return_value, float x, float y) { // Compute result and store in caller-provided slot *return_value x * y; // Signal that no Python exception occurred return 0; }注意实现中的两个关键点计算结果写入*return_value函数末尾返回0表示无异常。三、Python 侧声明cuda.declare_device()在 Python 中声明外部设备函数使用cuda.declare_device(name, sig)name外部函数的名称字符串。sig该函数的 Numba 签名例如float32(float32, float32)。返回的描述符descriptor名称不必与外部函数名一致。例如mul cuda.declare_device(mul_f32_f32, float32(float32, float32))声明之后在 kernel 内调用mul(a, b)会被翻译为对编译代码中mul_f32_f32(a, b)的调用。从源码看numba/cuda/decorators.py 中declare_device的实现要点def declare_device(name, sig): argtypes, restype sigutils.normalize_signature(sig) if restype is None: msg Return type must be provided for device declarations raise TypeError(msg) return declare_device_function(name, restype, argtypes)即声明时签名中的返回类型必须提供否则抛出TypeError。随后 numba/cuda/compiler.py 中的declare_device_function/declare_device_function_template会将外部函数注册到 CUDA target 的 typing context 和 target context 中sig typing.signature(restype, *argtypes) extfn ExternFunction(name, sig) # ... fndesc funcdesc.ExternalFunctionDescriptor( namename, restyperestype, argtypesargtypes) typingctx.insert_user_function(extfn, device_function_template) targetctx.insert_user_function(extfn, fndesc)这解释了为什么 kernel 内调用描述符时编译器能够直接生成对外部符号name的调用typing 阶段通过ExternFunction完成类型检查lowering 阶段通过ExternalFunctionDescriptor生成调用。四、传递指针让 C 函数与 Python kernel 的调用约定对齐Numba 的调用约定calling convention要求数组参数传递多个值数据指针、shape、stride 以及其他元信息。这与绝大多数只期望一个数据指针的 C/C 函数不兼容。因此为了让 C 设备代码与 Python kernel 的调用约定对齐数组参数必须用 C 指针类型来声明。例如假设外部函数原型为摘自 functions.cu 中的sum_reduceextern C __device__ int sum_reduce( float* return_value, float* array, int n );在 Python 中声明如下signature float32(CPointer(float32), int32) sum_reduce cuda.declare_device(sum_reduce, signature)其中CPointer(float32)即 Numba 的 C 指针类型。test_ex_from_buffer测试用例numba/cuda/tests/doc_examples/test_ffi.py中给出了完整用法。使用cffi.FFI.from_buffer()获取数组数据指针要在 kernel 内获得指向数组数据的指针传给外部函数使用cffi.FFI实例的from_buffer()方法。示例import cffi ffi cffi.FFI() cuda.jit(link[functions_cu]) def reduction_caller(result, array): array_ptr ffi.from_buffer(array) result[()] sum_reduce(array_ptr, len(array))其中result和array都是float32数据的数组在测试中x np.arange(10).astype(np.float32)r np.ndarray((), dtypenp.float32)。array_ptr是在 device 上指向数组数据的指针可直接作为sum_reduce的第二个参数传入。关于 cffi 依赖该示例依赖cffi包。测试类上使用了skip_unless_cffi装饰器见 test_ffi.py说明没有安装cffi时相关测试会被跳过实际使用ffi.from_buffer()前需要确保环境中有cffi。五、链接与调用cuda.jit(link[...])详解cuda.jit装饰器的link关键字参数接受一个文件名列表文件名可以是绝对路径也可以是相对于当前工作目录的路径numba/cuda/decorators.py。文件的处理规则见 numba/cuda/cudadrv/driver.py 的add_file_guess_ext以.cu结尾的文件使用NVRTCNVIDIA Runtime Compiler编译为 PTX再链接进 kernel其他文件直接传给 CUDA Linker按其扩展名推断类型如.ptx、.cubin、.fatbin等无扩展名或未知扩展名会抛出RuntimeError。驱动侧Linker.add_cu的实现driver.py展示了.cu文件的完整处理链def add_cu(self, cu, name): with driver.get_active_context() as ac: dev driver.get_device(ac.devnum) cc dev.compute_capability ptx, log nvrtc.compile(cu, name, cc) # ... 可选地输出汇编config.DUMP_ASSEMBLY # Link the programs PTX using the normal linker mechanism ptx_name os.path.splitext(name)[0] .ptx self.add_ptx(ptx.encode(), ptx_name)即.cu源码先用 NVRTC 针对当前设备计算能力compute capability编译为 PTX再走与.ptx相同的 PTX 链接路径。NVRTC 的 Python 封装位于 numba/cuda/cudadrv/nvrtc.py编译失败时会把 NVRTC 编译日志包含在NVRTC Compilation failure异常中便于排查。下面是在 kernel 中调用上文声明的mul()其实现mul_f32_f32位于文件functions.cu的完整示例from numba import cuda import numpy as np import os # Declaration of the foreign function mul cuda.declare_device(mul_f32_f32, float32(float32, float32)) # Path to the source containing the foreign function basedir os.path.dirname(os.path.abspath(__file__)) functions_cu os.path.join(basedir, ffi, functions.cu) # Kernel that links in functions.cu and calls mul cuda.jit(link[functions_cu]) def multiply_vectors(r, x, y): i cuda.grid(1) if i len(r): r[i] mul(x[i], y[i]) # Generate random data N 32 np.random.seed(1) x np.random.rand(N).astype(np.float32) y np.random.rand(N).astype(np.float32) r np.zeros_like(x) # Run the kernel multiply_vectors1, 32 # Sanity check - ensure the results match those expected np.testing.assert_array_equal(r, x * y)这段代码就是仓库中test_ex_linking_cu测试numba/cuda/tests/doc_examples/test_ffi.py的主体可直接照抄运行。使用限制link与 CUDA 模拟器cudasim不兼容decorators.py中当启用config.ENABLE_CUDASIM且提供了link时会抛出NotImplementedError“Cannot link PTX in the simulator”link对设备函数deviceTrue无效会抛出ValueError文档还提醒上面的最小示例主要为了演示外部函数调用由于 grid 很小且外部函数负载很轻不期望它有特别出色的性能。六、CUDA C/C 支持的前提条件使用 NVRTC 编译和链接 CUDA C/C 代码时需要注意以下前提见 cuda_ffi.rst 的 C/C Support 一节仅在使用 NVIDIA Bindings 时可用通过环境变量NUMBA_CUDA_USE_NVIDIA_BINDING控制numba/core/config.py 中默认值为 0结合 numba/cuda/cudadrv/driver.py 中Linker.new()对USE_NV_BINDING的分支选择可确认该开关会影响链接器的选择。需要与已安装的 NVIDIA CUDA Bindings 版本匹配的 NVRTC 库。NVRTC 的加载通过open_cudalib(nvrtc)完成numba/cuda/cudadrv/nvrtc.py加载失败会抛出NvrtcSupportError。CUDA include 路径默认在 Linux 为/usr/local/cuda/includeWindows 为$env:CUDA_PATH\include可通过环境变量NUMBA_CUDA_INCLUDE_PATH修改。配置实现见 numba/core/config.py该变量未设置时使用平台默认路径。include 目录只会以-I{config.CUDA_INCLUDE_PATH}的形式提供给 NVRTCnumba/cuda/cudadrv/nvrtc.py不支持下额外 include。也就是说.cu文件中只能依赖 CUDA 自带的头文件。环境变量速查环境变量作用默认值NUMBA_CUDA_USE_NVIDIA_BINDING是否使用 NVIDIA Bindings影响链接器与 NVRTC 支持0默认关闭NUMBA_CUDA_INCLUDE_PATHCUDA 头文件目录NVRTC 的 include 路径Linux/usr/local/cuda/includeWindows$env:CUDA_PATH\include七、完整示例汇总外部函数CUDA C保存为functions.cuextern C __device__ int mul_f32_f32( float* return_value, float x, float y) { *return_value x * y; return 0; }Python 调用端from numba import cuda import numpy as np mul cuda.declare_device(mul_f32_f32, float32(float32, float32)) cuda.jit(link[functions.cu]) def multiply_vectors(r, x, y): i cuda.grid(1) if i len(r): r[i] mul(x[i], y[i]) N 32 x np.random.rand(N).astype(np.float32) y np.random.rand(N).astype(np.float32) r np.zeros_like(x) multiply_vectors1, 32 np.testing.assert_array_equal(r, x * y)注意link[functions.cu]中的路径为相对当前工作目录的路径在仓库测试中文件路径是通过os.path.abspath构造的绝对路径实际使用时两者皆可。指针传递示例ffi.from_bufferimport cffi from numba import cuda import numpy as np signature float32(CPointer(float32), int32) sum_reduce cuda.declare_device(sum_reduce, signature) ffi cffi.FFI() cuda.jit(link[functions.cu]) def reduction_caller(result, array): array_ptr ffi.from_buffer(array) result[()] sum_reduce(array_ptr, len(array)) x np.arange(10).astype(np.float32) r np.ndarray((), dtypenp.float32) reduction_caller1, 1 np.testing.assert_allclose(np.sum(x), r[()])八、进一步阅读官方文档原文docs/source/cuda/cuda_ffi.rst文档配套的可运行测试numba/cuda/tests/doc_examples/test_ffi.py包含test_ex_linking_cu与test_ex_from_buffer两个用例与外部函数实现 numba/cuda/tests/doc_examples/ffi/functions.cu声明 API 实现numba/cuda/decorators.pydeclare_device与 numba/cuda/compiler.pydeclare_device_function链接器与 NVRTC 调用链numba/cuda/cudadrv/driver.pyLinker及add_cu_file/add_file_guess_ext、numba/cuda/cudadrv/nvrtc.pyNVRTC 链接相关测试numba/cuda/tests/cudadrv/test_linker.py环境变量配置numba/core/config.pyNUMBA_CUDA_USE_NVIDIA_BINDING、numba/core/config.pyNUMBA_CUDA_INCLUDE_PATH相关 CUDA 主题文档CUDA 编程模型概述、CUDA kernel 编写、NVIDIA Bindings 说明赞分享编译器高性能计算【免费下载链接】numbaNumPy aware dynamic Python compiler using LLVM项目地址https://gitcode.com/gh_mirrors/nu/numba点击查看免费下载相关推荐Numba CUDA 交叉编译指南使用 cuda.compile 将 Python 函数编译为 PTX / LTO-IR 并接入 C/C 程序Numba CUDA 交叉编译指南使用 cuda.compile 将 Python 函数编译为 PTX / LTO IR 并接入 C/C 程序 Numba编译器高性能计算Nintendo Switch大气层系统完整指南从零开始掌握自定义固件的终极教程Nintendo Switch大气层系统完整指南从零开始掌握自定义固件的终极教程 大气层系统Atmosphere 是Nintendo Switch最专业、编译器高性能计算Julia FFI编程终极指南无缝集成C库的完整教程Julia FFI编程终极指南无缝集成C库的完整教程 Julia作为一种高性能编程语言不仅自身拥有卓越的计算能力还通过外部函数接口FFI与C语言等低级编程语言编译器语言运行时标准库JIT编译上一篇终极Dokploy扩展开发完全指南自定义插件与AI功能扩展详解下一篇CPython xml.dom.pulldom 实战指南拉取式 XML 解析与局部 DOM 树的构建创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考