Warp 外部编译扩展 API 实战用wp.build_experimental注册 C 内置函数与原生值类型【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp本篇技术指南围绕 Warp 的warp.build_experimental模块展开它是 Warp 为外部 C/CUDA 库提供的实验性扩展接口用于把外部 C 函数与值类型接入 Warp 内核编译管线。在 Warp 中基于该模块你可以让生成的 kernel 直接调用第三方 C 设备函数、在 kernel 内构造与读写外部 C 结构体、把外部头文件与 include 路径注入模块级编译输入并在 AOT 编译时拿到产物路径。读完本文你将掌握wp.build_experimental.add_builtin、wp.build_experimental.add_native_type、wp.ModuleBuildOptions三个核心 API 的完整用法、底层编译校验机制与缓存失效规则并能够在自己基于 Warp 的库/插件中接入外部原生代码。1. 什么是warp.build_experimentalwarp.build_experimental是 Warp 面向基于 Warp 构建的包提供的扩展入口。这类包往往需要让外部 C 声明对生成的 kernel 可见、调用外部 C 函数、与库自定义的值类型交换数据以及在另一个运行时里加载生成的代码。在没有该模块之前这些需求只能通过给 Warp 打补丁或直接 importwarp._src内部实现来完成。该模块的 API 参考位于 docs/api_reference/warp_build_experimental.rst其公开成员只有两个——add_builtin与add_native_type它们由 warp/build_experimental.py 从warp._src.external_build重导出from warp._src.external_build import add_builtin as add_builtin from warp._src.external_build import add_native_type as add_native_type __all__ [add_builtin, add_native_type]需要特别注意的是整个模块的 API 均为实验性。模块 docstring 明确声明此模块的 API 是实验性的在外部编译特性的开发过程中可能不做弃用期直接变更见 warp/build_experimental.pyadd_builtin/add_native_type的文档字符串中也重复了这一警告。设计文档 design/external-compilation-extensions.md状态为 In Progress指出OptiX 是第一个消费者——一个 OptiX 插件需要其头文件在 NVRTC 编译期间可见、需要 OptiX 设备函数的可调用包装、OptiX 兼容的程序名与启动参数、以及生成的 PTX 路径而这些需求本身并不带 OptiX 特性因此该 API 被设计为对任意 C/CUDA 库通用。设计上的核心分工是Warp 继续拥有源码生成、哈希、编译与产物命名的所有权插件只提供其模块所需的额外声明与 ABI 契约。这不是一个通用的 C 构建系统外部实现代码必须通过被 include 的头文件可见编译/链接任意外部翻译单元或库被明确列为非目标。2.add_builtin把外部 C 函数注册为 kernel 可调用的内置函数wp.build_experimental.add_builtin()是 Warp 内部内置函数注册机制的一个窄封装只暴露具有稳定外部含义的参数。其完整签名见 warp/_src/external_build.pywp.build_experimental.add_builtin( name: str, input_types: Mapping[str, type] | None None, value_type: type | None None, *, native_name: str | None None, doc: str , ) - wp.Function2.1 参数语义参数含义说明nameWarp kernel 中调用该函数所用的名字注册后在 kernel 里以wp.name形式调用必须是不含特殊字符的 Python 标识符否则抛出ValueErrorinput_types从参数名到 Warp 类型的有序映射该映射定义恰好一个重载参数顺序是有意义的参数名须为合法 Python 标识符value_typeWarp 返回值类型传None表示函数返回voidnative_name精确限定的 C 函数名省略时默认为fwp::{name}必须匹配^(?:::)?[A-Za-z_]\w*(?:::[A-Za-z_]\w*)*$即合法的限定 C 标识符否则抛ValueErrordoc函数描述必须是str注册后函数在 kernel 内通过wp.name调用但不允许从 Python 直接调用。返回值为包含该注册重载的规范wp.Function对象。native_name会被拆分为命名空间与函数名两部分见_split_native_namewarp/_src/external_build.py例如native_nameexternal::square会得到 namespaceexternal::与函数名square以::开头的全局限定名如::foo会被识别为全局命名空间限定。测试 warp/tests/test_external_build.py 验证了这一点注册后overload.namespace external::、overload.native_func square。2.2 注册属性不可微、隐藏、不导出外部内置函数以固定的内部参数注册见 warp/_src/external_build.pyexportFalse、hiddenTrue、is_differentiableFalse。这意味着不可微不参与 Warp 的自动微分也没有梯度钩子隐藏不会出现在静态存根.pyi或生成的 builtin 参考文档中不导出仅能在 kernel 代码里调用。由于所有注册共享 Warp 全局的内置函数命名空间插件之间可能撞名因此设计建议插件为自己的名字加前缀。向已有的 Warp 操作追加重载是允许的有时甚至是刻意的。2.3 幂等注册与冲突拒绝注册契约由签名有序输入类型 返回值类型 命名空间 原生函数名共同构成。重复注册相同签名、相同结果类型、相同原生目标时add_builtin直接返回已存在的Function不会报错但复用同一签名搭配不同返回值类型或不同原生目标时抛出RuntimeError(Cannot register conflicting external builtin overload ...)见 warp/_src/external_build.py。测试 warp/tests/test_external_build.py 还验证了一个细节相同签名注册两次返回同一个函数对象assertIs(duplicate, function)而value_type改为wp.int32后立刻冲突报错。2.4 完整示例注册并调用外部 C 函数参考测试文件 warp/tests/test_external_build.py 与设计文档中的示例一个典型的外部函数注册如下import ctypes import warp as wp wp.build_experimental.add_builtin( my_addon_scale_color, # kernel 中调用为 wp.my_addon_scale_color(...) {value: NativeColor, factor: wp.float32}, # 参数顺序即 C 形参顺序 NativeColor, # 返回值类型 native_namemy_addon::scale_color, # 精确限定 C 函数名 docScale a color by a factor., )对应的 C 声明需要提前通过模块 preamble/include 注入见第 4 节例如namespace my_addon { struct Color { float r, g, b; }; CUDA_CALLABLE inline Color scale_color(Color value, float factor) { return Color{value.r * factor, value.g * factor, value.b * factor}; } }之后便可在 kernel 中直接调用wp.kernel def use_kernel(pixels: wp.array[NativeColor]): i wp.tid() pixels[i] wp.my_addon_scale_color(pixels[i], 2.0)测试 warp/tests/test_external_build.py 展示了完整的运行验证colors[i] wp._test_native_scale(pixel.color, 2.0)之后colors.numpy()的各通道值[0.0, 2.0]、[1.0, 1.0]、[6.0, 6.0]与期望完全一致。3.add_native_type把外部 C 值类型注册为 Warp 注解与数组 dtypewp.build_experimental.add_native_type()注册一个ctypes.Structure子类作为外部 C 值类型的宿主 ABI 描述。其签名见 warp/_src/external_build.pywp.build_experimental.add_native_type( ctype: type[ctypes.Structure], *, native_name: str, fields: Mapping[str, type] | None None, initializer: str | None None, ) - type[ctypes.Structure]同一个 ctypes 类身兼三职既是宿主hostABI 声明又是 Warp kernel 的参数注解还是wp.array的dtype。这避免了维护一套平行的包装类型层级也让 ctypes 与 NumPy 互操作保持直接。3.1 参数语义参数含义说明ctypectypes.Structure子类必须与 C 类型具有相同的 size、alignment 与字段布局Warp 在编译期校验 size 与 alignmentnative_name精确限定的 C 类型名规则与add_builtin相同必须为合法的限定 C 标识符fields从公开 C 数据成员名到 Warp 类型的有序映射传None表示不透明opaque类型映射中每个字段必须存在于_fields_中initializer取值None或aggregateaggregate允许在 kernel 内用全部字段做聚合构造此时fields必须按声明顺序列出 ctypes 的每一个字段C 侧的类型必须是标准布局standard-layout且可平凡复制trivially copyable。Warp 只按字节复制值不拥有值引用的资源也绝不会调用外部析构函数。3.2 不透明类型opaque与字段暴露类型fieldsNone不透明值可以跨 kernel 与 builtin 边界传递、放进 Warp struct、存进数组但 kernel 代码无法读取其成员。测试中的NativeImage即为此类——只含一个handle: c_uint64注册时未给fields。fields{...}暴露成员把公开 C 数据成员暴露为 Warp 类型。默认构造总是被允许的initializeraggregate额外允许从字段值构造。由于 C17 聚合初始化是按位置进行的该模式要求fields按声明顺序覆盖全部 ctypes 字段否则可能静默地把值初始化到错误的 C 成员上——违反此约束会抛ValueError见 warp/_src/external_build.py。原生类型不会自动获得运算符或微分能力对原生数组使用requires_gradTrue会被拒绝测试 warp/tests/test_external_build.py 验证了ValueError: automatic differentiation。所有针对原生类型的运算都需要由插件以 builtin 形式另行注册。3.3 注册期校验Python 侧注册时会逐字段校验warp/_src/external_build.py每个暴露字段必须存在于ctype._fields_中字段的 ctypes 存储大小必须与所映射 Warp 类型的大小一致如c_int16映射wp.int32会报ValueError: ...ctypes storage..._fields_中长度为 2 以外的项位域布局被直接拒绝initializeraggregate要求暴露字段名集合与声明字段名集合完全一致且顺序一致。测试 warp/tests/test_external_build.py 对上述规则一一做了断言包括字段存储大小不匹配、aggregate 缺 fields、字段重排、字段缺失四种失败场景。3.4 编译期 ABI 校验C 侧ctypes 定义是对 C 布局的承诺因此 Warp 在两端都做校验注册期Python校验存储大小编译期对每个引用了该类型的模块在生成源码中发射static_assert。CPU 与 CUDA 后端校验的属性互补见 design/external-compilation-extensions.md契约CPUClangCUDANVRTC可平凡复制且标准布局static_assertstatic_assert类型 size 与 alignmentstatic_assertstatic_assert暴露字段类型decltype/ 相等性 traitdecltype/ 相等性 trait暴露字段偏移__builtin_offsetof不检查暴露字段大小不检查sizeof(((T*)0)-field)两个后端检查互补属性的原因是 NVRTC 在受支持的配置下不能稳定提供可用的offsetof表达式而指针式sizeof技巧在所有环境下都可用。因此纯 CUDA 插件不会获得字段偏移校验需要自行保证原生成员顺序与填充一致——把同一类型在 CPU 编译一次是最便宜的获得该项检查的方式。测试 warp/tests/test_external_build.py 验证了生成源码中确实包含wp_external_type_is_samedecltype(((warp_test::Color*)0)-r), wp::float32::value之类的契约断言。3.5 类型码schema与幂等注册注册的 schema 包含限定 C 名、size、alignment、初始化策略、暴露字段的名字/类型/偏移。它被 SHA-256 哈希截断为一个短类型码ntdigest并参与模块哈希见 warp/_src/external_build.py。等效 schema 的重复注册会得到相同的稳定类型码因此注册是幂等的——测试用两个不同的 ctypes 类模拟模块 reload注册相同 schema 均成功而注册不同 schema 到同一 C 名字则抛RuntimeError: ...different definition...warp/tests/test_external_build.py。需要澄清的是幂等性并不等于 reload 安全。一个新建的 ctypes 类在重载解析与参数匹配时与原始类不可互换reload 安全的类型匹配属于未来工作。此外元数据以_wp_native_type_/_wp_native_vars_私有属性挂在 ctypes 类上warp/_src/external_build.py这是为了让既有的面向 struct 的代码生成保持通用但它不是扩展协议——插件应当只使用wp.build_experimental。3.6 完整示例Color / Pixel / Image综合测试文件顶部的注册代码warp/tests/test_external_build.pyclass NativeColor(ctypes.Structure): _fields_ [(r, ctypes.c_float), (g, ctypes.c_float), (b, ctypes.c_float)] class NativePixel(ctypes.Structure): _fields_ [(color, NativeColor), (index, ctypes.c_int32)] class NativeImage(ctypes.Structure): _fields_ [(handle, ctypes.c_uint64)] wp.build_experimental.add_native_type( NativeColor, native_namewarp_test::Color, fields{r: wp.float32, g: wp.float32, b: wp.float32}, initializeraggregate, ) wp.build_experimental.add_native_type( NativePixel, native_namewarp_test::Pixel, fields{color: NativeColor, index: wp.int32}, # 字段可嵌套另一个原生类型 initializeraggregate, ) wp.build_experimental.add_native_type(NativeImage, native_namewarp_test::Image) # 不透明类型注册后这些类型可以在 kernel 里构造NativeColor(wp.float32(i), const.g, 3.0)、作为 kernel 参数、作为 Warp struct 字段、作为wp.array的 dtype并能与 NumPy 完整往返——包括带填充padded的类型测试 warp/tests/test_external_build.py 用c_uint8c_uint32的组合验证了wp.array(...).numpy()输出可以重新喂回wp.array。不透明类型的 NumPy 往返 dtype 为np.dtype(fV{ctypes.sizeof(NativeImage)})且拒绝结构化structuredNumPy 数据作为其数组输入warp/tests/test_external_build.py。4.ModuleBuildOptions模块级编译输入include 目录、preamble、依赖文件外部类型与函数要真正可编译必须把对应的头文件、include 路径和依赖文件交给编译器。这部分由wp.ModuleBuildOptions承担实现位于 warp/_src/context.py。wp.ModuleBuildOptions( *, extra_cuda_include_dirs: Sequence[str | os.PathLike[str]] | None None, extra_cpu_include_dirs: Sequence[str | os.PathLike[str]] | None None, extra_cuda_preamble: str , extra_cpu_preamble: str , extra_build_dependencies: Sequence[str | os.PathLike[str]] | None None, ) - ModuleBuildOptions4.1 五个参数的作用参数作用extra_cuda_include_dirs仅 CUDANVRTC编译时追加的 include 目录每项必须是存在的绝对路径extra_cpu_include_dirs仅 CPUClang编译时追加的 include 目录每项必须是存在的绝对路径extra_cuda_preamble注入 CUDA 源码的额外文本插在 Warp 头文件之后、codegen 专用 cast 宏与生成代码之前extra_cpu_preamble同上针对 CPU 编译extra_build_dependencies内容被哈希进模块哈希的文件列表每项必须是存在的绝对路径include 目录与依赖文件必须是绝对路径且必须存在但构造函数与merged()不检查这一点——校验推迟到模块选项被解析用于构建时。这是有意设计保持组合操作廉价把错误抛在缺失路径真正造成影响的地方。4.2 preamble 的插入位置为什么重要preamble 被插入在 Warp 原生头文件之后、codegen-only cast 宏与生成代码之前。设计文档 design/external-compilation-extensions.md 记录了这段演进最初选择把 preamble 放在最前面但结果是错的——CPU 上 Clang 会把预编译的builtin.h注入到翻译单元最前导致靠前的 preamble 在 CPU 上落在 Warp 头文件之后、在 NVRTC 下却落在之前而且 CPU 行为还随warp.config.use_precompiled_headers翻转。放在原生头文件之后两个后端行为才一致与 PCH 状态无关。这一位置带来的实际收益测试 warp/tests/test_external_build.py 验证外部头文件可以使用 Warp 的公开宏如CUDA_CALLABLE普通的 C 函数式 cast如float(x)、int(x)不会被 codegen 的 cast 宏改写保持为普通 C生成的 kernel 能看到 preamble 声明的全部内容。代价是preamble不能定义被 Warp 自身头文件消费的宏。4.3merged()非破坏性组合merged(*others)返回一个新对象不修改任何输入。组合语义warp/_src/context.pyinclude 目录与依赖文件保留首次出现去重preamble按参数顺序拼接非空值之间以换行分隔。测试 warp/tests/test_external_build.py 验证了base.merged(addon)的完整结果例如 CUDA include 目录[cuda/base, shared, cuda/addon]shared只出现一次、preamble#define BASE_CUDA 1\n#define ADDON_CUDA 1\n且base本身保持不变。这使得独立开发的插件可以安全组合不会互相覆盖或复制全部选项。4.4 如何应用到模块set_module_optionsModuleBuildOptions通过既有模块选项 API 的新键extra_build_options应用该键在 warp/_src/context.py 被解析并归一化wp.set_module_options( {extra_build_options: wp.ModuleBuildOptions( extra_cuda_include_dirs[include_dir], extra_cuda_preamble#include my_addon.h, extra_build_dependencies[header_path], )}, modulesome_module, # 也接受 Module 对象或模块名 )set_module_options的module参数现在接受Module对象或模块名不再局限于 Python 模块。插件典型的使用模式design/external-compilation-extensions.md在 import 时注册语言元素类型与函数再把构建选项合并进每个消费模块的extra_build_options且必须在模块首次 JIT 启动或 AOT 编译之前完成。4.5 哈希与缓存失效规则模块哈希覆盖解析后的路径、preamble 文本、依赖文件内容、代码生成选项、以及被引用的 Warp 可见扩展契约外部 builtin 身份与原生类型码。两点关键行为Warp不会扫描传递的 C include——任何内容变化都应使缓存失效的头文件必须由插件显式列入extra_build_dependencies依赖内容在模块哈希被重新计算时wp.set_module_options()之后或新进程才被重读而不是每次 launch 都读——因此运行中的进程里修改头文件不会重建已加载模块。此外就地修改ModuleBuildOptions实例不会使已编译模块失效调用者必须通过wp.set_module_options()重新应用这是唯一的显式失效点warp/_src/context.py 与 design/external-compilation-extensions.md。测试 warp/tests/test_external_build.py 验证了依赖文件内容从#define EXTERNAL_VALUE 1改为2后模块哈希确实不同。5. 完整串联示例一个最小外部插件综合第 2~4 节一个最小插件的完整形态如下参考 warp/tests/test_external_build.pyimport ctypes import warp as wp # 1) 定义宿主 ABI 描述 class Color(ctypes.Structure): _fields_ [(r, ctypes.c_float), (g, ctypes.c_float), (b, ctypes.c_float)] # 2) import 时注册语言元素 wp.build_experimental.add_native_type( Color, native_namemy_addon::Color, fields{r: wp.float32, g: wp.float32, b: wp.float32}, initializeraggregate, ) wp.build_experimental.add_builtin( my_addon_scale, {value: Color, factor: wp.float32}, Color, native_namemy_addon::scale, ) # 3) 构造并应用模块级构建选项 PREAMBLE namespace my_addon { struct Color { float r, g, b; }; CUDA_CALLABLE inline Color scale(Color value, float factor) { return Color{value.r * factor, value.g * factor, value.b * factor}; } } wp.set_module_options( { extra_build_options: wp.ModuleBuildOptions( extra_cpu_preamblePREAMBLE, extra_cuda_preamblePREAMBLE, ), }, modulemy_app_module, # 消费该插件的模块 ) # 4) 在 kernel 中使用 wp.kernel(modulemy_app_module) def scale_kernel(colors: wp.array[Color], factor: wp.float32): i wp.tid() colors[i] wp.my_addon_scale(colors[i], factor)启动后即可wp.launch(scale_kernel, ...)正常执行。测试中还展示了原生值作为 kernel 参数时wp.Tape的反向传播不受影响不透明值作为普通参数参与梯度计算见test_native_value_tape_adjointwarp/tests/test_external_build.py。6. 扩展 ABI 与 AOT 产物路径warp.build_experimental之外外部编译特性还涉及两处配套 API 变化详见 design/external-compilation-extensions.md6.1 可选入口点 ABIexternal_constant_paramswp.kernel(entry_point_abi...)新增可选参数默认warp行为完全不变另一个取值是external_constant_params。该 ABI 生成一个无参数的extern CCUDA 入口点其唯一的 Warp struct 参数从模块级常量内存符号params中读取extern C { __constant__ __align__(alignof(MyParams)) unsigned char params[sizeof(MyParams)]; } extern C __global__ void my_kernel() { MyParams var_p *reinterpret_castconst MyParams*(params); // ... kernel body ... }入口点使用 kernel 的 mangled 名不带_cuda_kernel_forward后缀配合strip_hashTrue时Warp 使用去掉哈希后缀的 kernel key给外部运行时一个稳定的设备侧符号供查找。该 ABI 被刻意约束使用它的 kernel 必须仅限 CUDA、恰好接受一个 Warp struct 参数、设置enable_backwardFalse、避免wp.tid()、共享内存 tile 与确定性原子操作且不能传给wp.launch()。由于params是单一模块级符号同一模块内所有external_constant_paramskernel必须使用同一个 struct 类型混用类型会在代码生成期报错。6.2compile_aot_module返回产物路径wp.compile_aot_module()从返回None改为按目标顺序先设备、后显式arch值返回pathlib.Path对象列表。消费者不再需要重建 Warp 私有的缓存与架构命名规则忽略旧None返回值的调用者不受影响。7. 设计取舍与替代方案设计文档 design/external-compilation-extensions.md 记录了四个被否决的替代方案理解它们有助于把握 API 边界直接暴露内部注册 API代码改动小但 API 承诺大——内部add_builtin的大多数参数只服务于 Warp 自身的内置库没有稳定的外部契约进程级全局编译器设置全局add_header()/add_include_directory()/ 宏 / C 标准设置原型方便但后续每个模块都会继承变更插件无法隔离冲突设置缓存失效语义含糊并发编译还会观察到无关状态。模块级选项保留了实用能力且使所有权与哈希显式化注册源文件或库需要跨越 NVRTC、嵌入式 Clang、NVJitLink、目标架构、可重定位设备代码与缓存产物的公开链接模型头文件可见的实现已覆盖首批消费者独立源码/库支持可后续再加而不破坏契约把原生类型建模为 Warp struct 或新包装对象Warp struct 是外部定义的复制品、无法证明其 ABI新公开包装类型会在注解、数组、常量与 NumPy 互操作上引发身份与转换问题。ctypes 类本身已提供宿主布局无需额外对用户可见的类型。8. 限制与未来工作当前实验性 API 面被限定为模块构建输入、外部注册、入口点 ABI、AOT 产物路径返回契约。已知限制包括design/external-compilation-extensions.md注册表是进程全局的没有 unregister 机制动态注册的内置函数名对静态类型检查器不可见模块 reload 后新建的 ctypes 类与原始类在重载解析/参数匹配上不可互换不扫描传递 C include依赖头文件必须显式列出。文档列出的可行后续方向链接翻译单元、更强的 CUDA ABI 检查、外部 builtin 的梯度契约、更多构造策略、更多命名入口点 ABI、显式注册表生命周期——这些都应当跟随真实消费者出现而不是提前暴露编译器内部。9. 如何验证与测试仓库内的测试覆盖了本节 API 的全部行为可作为实现契约的权威参照warp/tests/test_external_build.py注册校验、冲突与幂等重注册、merged()顺序与非破坏性、依赖文件与外部 builtin 契约哈希、双后端 preamble 位置以及不透明/字段描述原生值作为 kernel 参数、builtin 结果、struct 字段与数组 dtype 的 NumPy 往返warp/tests/test_codegen.pypreamble 与后端专属 include 目录哈希、external_constant_params代码生成及其校验错误warp/tests/test_compilation.py 与 warp/tests/aot/test_module_aot.pyinclude 目录到达 Clang、CPU/CUDA 设备与架构列表的 AOT 产物路径warp/tests/cuda/test_occupancy.py 与 warp/tests/cuda/test_kernel_attributes.py外部入口点的占用率与 kernel 属性查询warp/tests/test_modules_lite.py确认内部add_builtin不属于公开 API。运行方式与仓库其他测试一致例如在仓库根目录执行python -m warp.tests.test_external_build或经由warp/tests/__main__.py的测试套件入口。设计文档同时提到针对真实外部消费者otk-pyoptix不 import 任何warp._src的验证目前为手动方式。最后提醒wp.build_experimental的experimental前缀是刻意的它让临时状态在 import 处可见。如果这些 API 未来达到 Warp 的稳定性与支持标准可能像warp.jax_experimental被提升为warp.jax一样被提升到稳定的warp.build命名空间——任何提升都会制定自己的兼容性与弃用计划。【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Warp 外部编译扩展 API 实战：用 `wp.build_experimental` 注册 C++ 内置函数与原生值类型
Warp 外部编译扩展 API 实战用wp.build_experimental注册 C 内置函数与原生值类型【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp本篇技术指南围绕 Warp 的warp.build_experimental模块展开它是 Warp 为外部 C/CUDA 库提供的实验性扩展接口用于把外部 C 函数与值类型接入 Warp 内核编译管线。在 Warp 中基于该模块你可以让生成的 kernel 直接调用第三方 C 设备函数、在 kernel 内构造与读写外部 C 结构体、把外部头文件与 include 路径注入模块级编译输入并在 AOT 编译时拿到产物路径。读完本文你将掌握wp.build_experimental.add_builtin、wp.build_experimental.add_native_type、wp.ModuleBuildOptions三个核心 API 的完整用法、底层编译校验机制与缓存失效规则并能够在自己基于 Warp 的库/插件中接入外部原生代码。1. 什么是warp.build_experimentalwarp.build_experimental是 Warp 面向基于 Warp 构建的包提供的扩展入口。这类包往往需要让外部 C 声明对生成的 kernel 可见、调用外部 C 函数、与库自定义的值类型交换数据以及在另一个运行时里加载生成的代码。在没有该模块之前这些需求只能通过给 Warp 打补丁或直接 importwarp._src内部实现来完成。该模块的 API 参考位于 docs/api_reference/warp_build_experimental.rst其公开成员只有两个——add_builtin与add_native_type它们由 warp/build_experimental.py 从warp._src.external_build重导出from warp._src.external_build import add_builtin as add_builtin from warp._src.external_build import add_native_type as add_native_type __all__ [add_builtin, add_native_type]需要特别注意的是整个模块的 API 均为实验性。模块 docstring 明确声明此模块的 API 是实验性的在外部编译特性的开发过程中可能不做弃用期直接变更见 warp/build_experimental.pyadd_builtin/add_native_type的文档字符串中也重复了这一警告。设计文档 design/external-compilation-extensions.md状态为 In Progress指出OptiX 是第一个消费者——一个 OptiX 插件需要其头文件在 NVRTC 编译期间可见、需要 OptiX 设备函数的可调用包装、OptiX 兼容的程序名与启动参数、以及生成的 PTX 路径而这些需求本身并不带 OptiX 特性因此该 API 被设计为对任意 C/CUDA 库通用。设计上的核心分工是Warp 继续拥有源码生成、哈希、编译与产物命名的所有权插件只提供其模块所需的额外声明与 ABI 契约。这不是一个通用的 C 构建系统外部实现代码必须通过被 include 的头文件可见编译/链接任意外部翻译单元或库被明确列为非目标。2.add_builtin把外部 C 函数注册为 kernel 可调用的内置函数wp.build_experimental.add_builtin()是 Warp 内部内置函数注册机制的一个窄封装只暴露具有稳定外部含义的参数。其完整签名见 warp/_src/external_build.pywp.build_experimental.add_builtin( name: str, input_types: Mapping[str, type] | None None, value_type: type | None None, *, native_name: str | None None, doc: str , ) - wp.Function2.1 参数语义参数含义说明nameWarp kernel 中调用该函数所用的名字注册后在 kernel 里以wp.name形式调用必须是不含特殊字符的 Python 标识符否则抛出ValueErrorinput_types从参数名到 Warp 类型的有序映射该映射定义恰好一个重载参数顺序是有意义的参数名须为合法 Python 标识符value_typeWarp 返回值类型传None表示函数返回voidnative_name精确限定的 C 函数名省略时默认为fwp::{name}必须匹配^(?:::)?[A-Za-z_]\w*(?:::[A-Za-z_]\w*)*$即合法的限定 C 标识符否则抛ValueErrordoc函数描述必须是str注册后函数在 kernel 内通过wp.name调用但不允许从 Python 直接调用。返回值为包含该注册重载的规范wp.Function对象。native_name会被拆分为命名空间与函数名两部分见_split_native_namewarp/_src/external_build.py例如native_nameexternal::square会得到 namespaceexternal::与函数名square以::开头的全局限定名如::foo会被识别为全局命名空间限定。测试 warp/tests/test_external_build.py 验证了这一点注册后overload.namespace external::、overload.native_func square。2.2 注册属性不可微、隐藏、不导出外部内置函数以固定的内部参数注册见 warp/_src/external_build.pyexportFalse、hiddenTrue、is_differentiableFalse。这意味着不可微不参与 Warp 的自动微分也没有梯度钩子隐藏不会出现在静态存根.pyi或生成的 builtin 参考文档中不导出仅能在 kernel 代码里调用。由于所有注册共享 Warp 全局的内置函数命名空间插件之间可能撞名因此设计建议插件为自己的名字加前缀。向已有的 Warp 操作追加重载是允许的有时甚至是刻意的。2.3 幂等注册与冲突拒绝注册契约由签名有序输入类型 返回值类型 命名空间 原生函数名共同构成。重复注册相同签名、相同结果类型、相同原生目标时add_builtin直接返回已存在的Function不会报错但复用同一签名搭配不同返回值类型或不同原生目标时抛出RuntimeError(Cannot register conflicting external builtin overload ...)见 warp/_src/external_build.py。测试 warp/tests/test_external_build.py 还验证了一个细节相同签名注册两次返回同一个函数对象assertIs(duplicate, function)而value_type改为wp.int32后立刻冲突报错。2.4 完整示例注册并调用外部 C 函数参考测试文件 warp/tests/test_external_build.py 与设计文档中的示例一个典型的外部函数注册如下import ctypes import warp as wp wp.build_experimental.add_builtin( my_addon_scale_color, # kernel 中调用为 wp.my_addon_scale_color(...) {value: NativeColor, factor: wp.float32}, # 参数顺序即 C 形参顺序 NativeColor, # 返回值类型 native_namemy_addon::scale_color, # 精确限定 C 函数名 docScale a color by a factor., )对应的 C 声明需要提前通过模块 preamble/include 注入见第 4 节例如namespace my_addon { struct Color { float r, g, b; }; CUDA_CALLABLE inline Color scale_color(Color value, float factor) { return Color{value.r * factor, value.g * factor, value.b * factor}; } }之后便可在 kernel 中直接调用wp.kernel def use_kernel(pixels: wp.array[NativeColor]): i wp.tid() pixels[i] wp.my_addon_scale_color(pixels[i], 2.0)测试 warp/tests/test_external_build.py 展示了完整的运行验证colors[i] wp._test_native_scale(pixel.color, 2.0)之后colors.numpy()的各通道值[0.0, 2.0]、[1.0, 1.0]、[6.0, 6.0]与期望完全一致。3.add_native_type把外部 C 值类型注册为 Warp 注解与数组 dtypewp.build_experimental.add_native_type()注册一个ctypes.Structure子类作为外部 C 值类型的宿主 ABI 描述。其签名见 warp/_src/external_build.pywp.build_experimental.add_native_type( ctype: type[ctypes.Structure], *, native_name: str, fields: Mapping[str, type] | None None, initializer: str | None None, ) - type[ctypes.Structure]同一个 ctypes 类身兼三职既是宿主hostABI 声明又是 Warp kernel 的参数注解还是wp.array的dtype。这避免了维护一套平行的包装类型层级也让 ctypes 与 NumPy 互操作保持直接。3.1 参数语义参数含义说明ctypectypes.Structure子类必须与 C 类型具有相同的 size、alignment 与字段布局Warp 在编译期校验 size 与 alignmentnative_name精确限定的 C 类型名规则与add_builtin相同必须为合法的限定 C 标识符fields从公开 C 数据成员名到 Warp 类型的有序映射传None表示不透明opaque类型映射中每个字段必须存在于_fields_中initializer取值None或aggregateaggregate允许在 kernel 内用全部字段做聚合构造此时fields必须按声明顺序列出 ctypes 的每一个字段C 侧的类型必须是标准布局standard-layout且可平凡复制trivially copyable。Warp 只按字节复制值不拥有值引用的资源也绝不会调用外部析构函数。3.2 不透明类型opaque与字段暴露类型fieldsNone不透明值可以跨 kernel 与 builtin 边界传递、放进 Warp struct、存进数组但 kernel 代码无法读取其成员。测试中的NativeImage即为此类——只含一个handle: c_uint64注册时未给fields。fields{...}暴露成员把公开 C 数据成员暴露为 Warp 类型。默认构造总是被允许的initializeraggregate额外允许从字段值构造。由于 C17 聚合初始化是按位置进行的该模式要求fields按声明顺序覆盖全部 ctypes 字段否则可能静默地把值初始化到错误的 C 成员上——违反此约束会抛ValueError见 warp/_src/external_build.py。原生类型不会自动获得运算符或微分能力对原生数组使用requires_gradTrue会被拒绝测试 warp/tests/test_external_build.py 验证了ValueError: automatic differentiation。所有针对原生类型的运算都需要由插件以 builtin 形式另行注册。3.3 注册期校验Python 侧注册时会逐字段校验warp/_src/external_build.py每个暴露字段必须存在于ctype._fields_中字段的 ctypes 存储大小必须与所映射 Warp 类型的大小一致如c_int16映射wp.int32会报ValueError: ...ctypes storage..._fields_中长度为 2 以外的项位域布局被直接拒绝initializeraggregate要求暴露字段名集合与声明字段名集合完全一致且顺序一致。测试 warp/tests/test_external_build.py 对上述规则一一做了断言包括字段存储大小不匹配、aggregate 缺 fields、字段重排、字段缺失四种失败场景。3.4 编译期 ABI 校验C 侧ctypes 定义是对 C 布局的承诺因此 Warp 在两端都做校验注册期Python校验存储大小编译期对每个引用了该类型的模块在生成源码中发射static_assert。CPU 与 CUDA 后端校验的属性互补见 design/external-compilation-extensions.md契约CPUClangCUDANVRTC可平凡复制且标准布局static_assertstatic_assert类型 size 与 alignmentstatic_assertstatic_assert暴露字段类型decltype/ 相等性 traitdecltype/ 相等性 trait暴露字段偏移__builtin_offsetof不检查暴露字段大小不检查sizeof(((T*)0)-field)两个后端检查互补属性的原因是 NVRTC 在受支持的配置下不能稳定提供可用的offsetof表达式而指针式sizeof技巧在所有环境下都可用。因此纯 CUDA 插件不会获得字段偏移校验需要自行保证原生成员顺序与填充一致——把同一类型在 CPU 编译一次是最便宜的获得该项检查的方式。测试 warp/tests/test_external_build.py 验证了生成源码中确实包含wp_external_type_is_samedecltype(((warp_test::Color*)0)-r), wp::float32::value之类的契约断言。3.5 类型码schema与幂等注册注册的 schema 包含限定 C 名、size、alignment、初始化策略、暴露字段的名字/类型/偏移。它被 SHA-256 哈希截断为一个短类型码ntdigest并参与模块哈希见 warp/_src/external_build.py。等效 schema 的重复注册会得到相同的稳定类型码因此注册是幂等的——测试用两个不同的 ctypes 类模拟模块 reload注册相同 schema 均成功而注册不同 schema 到同一 C 名字则抛RuntimeError: ...different definition...warp/tests/test_external_build.py。需要澄清的是幂等性并不等于 reload 安全。一个新建的 ctypes 类在重载解析与参数匹配时与原始类不可互换reload 安全的类型匹配属于未来工作。此外元数据以_wp_native_type_/_wp_native_vars_私有属性挂在 ctypes 类上warp/_src/external_build.py这是为了让既有的面向 struct 的代码生成保持通用但它不是扩展协议——插件应当只使用wp.build_experimental。3.6 完整示例Color / Pixel / Image综合测试文件顶部的注册代码warp/tests/test_external_build.pyclass NativeColor(ctypes.Structure): _fields_ [(r, ctypes.c_float), (g, ctypes.c_float), (b, ctypes.c_float)] class NativePixel(ctypes.Structure): _fields_ [(color, NativeColor), (index, ctypes.c_int32)] class NativeImage(ctypes.Structure): _fields_ [(handle, ctypes.c_uint64)] wp.build_experimental.add_native_type( NativeColor, native_namewarp_test::Color, fields{r: wp.float32, g: wp.float32, b: wp.float32}, initializeraggregate, ) wp.build_experimental.add_native_type( NativePixel, native_namewarp_test::Pixel, fields{color: NativeColor, index: wp.int32}, # 字段可嵌套另一个原生类型 initializeraggregate, ) wp.build_experimental.add_native_type(NativeImage, native_namewarp_test::Image) # 不透明类型注册后这些类型可以在 kernel 里构造NativeColor(wp.float32(i), const.g, 3.0)、作为 kernel 参数、作为 Warp struct 字段、作为wp.array的 dtype并能与 NumPy 完整往返——包括带填充padded的类型测试 warp/tests/test_external_build.py 用c_uint8c_uint32的组合验证了wp.array(...).numpy()输出可以重新喂回wp.array。不透明类型的 NumPy 往返 dtype 为np.dtype(fV{ctypes.sizeof(NativeImage)})且拒绝结构化structuredNumPy 数据作为其数组输入warp/tests/test_external_build.py。4.ModuleBuildOptions模块级编译输入include 目录、preamble、依赖文件外部类型与函数要真正可编译必须把对应的头文件、include 路径和依赖文件交给编译器。这部分由wp.ModuleBuildOptions承担实现位于 warp/_src/context.py。wp.ModuleBuildOptions( *, extra_cuda_include_dirs: Sequence[str | os.PathLike[str]] | None None, extra_cpu_include_dirs: Sequence[str | os.PathLike[str]] | None None, extra_cuda_preamble: str , extra_cpu_preamble: str , extra_build_dependencies: Sequence[str | os.PathLike[str]] | None None, ) - ModuleBuildOptions4.1 五个参数的作用参数作用extra_cuda_include_dirs仅 CUDANVRTC编译时追加的 include 目录每项必须是存在的绝对路径extra_cpu_include_dirs仅 CPUClang编译时追加的 include 目录每项必须是存在的绝对路径extra_cuda_preamble注入 CUDA 源码的额外文本插在 Warp 头文件之后、codegen 专用 cast 宏与生成代码之前extra_cpu_preamble同上针对 CPU 编译extra_build_dependencies内容被哈希进模块哈希的文件列表每项必须是存在的绝对路径include 目录与依赖文件必须是绝对路径且必须存在但构造函数与merged()不检查这一点——校验推迟到模块选项被解析用于构建时。这是有意设计保持组合操作廉价把错误抛在缺失路径真正造成影响的地方。4.2 preamble 的插入位置为什么重要preamble 被插入在 Warp 原生头文件之后、codegen-only cast 宏与生成代码之前。设计文档 design/external-compilation-extensions.md 记录了这段演进最初选择把 preamble 放在最前面但结果是错的——CPU 上 Clang 会把预编译的builtin.h注入到翻译单元最前导致靠前的 preamble 在 CPU 上落在 Warp 头文件之后、在 NVRTC 下却落在之前而且 CPU 行为还随warp.config.use_precompiled_headers翻转。放在原生头文件之后两个后端行为才一致与 PCH 状态无关。这一位置带来的实际收益测试 warp/tests/test_external_build.py 验证外部头文件可以使用 Warp 的公开宏如CUDA_CALLABLE普通的 C 函数式 cast如float(x)、int(x)不会被 codegen 的 cast 宏改写保持为普通 C生成的 kernel 能看到 preamble 声明的全部内容。代价是preamble不能定义被 Warp 自身头文件消费的宏。4.3merged()非破坏性组合merged(*others)返回一个新对象不修改任何输入。组合语义warp/_src/context.pyinclude 目录与依赖文件保留首次出现去重preamble按参数顺序拼接非空值之间以换行分隔。测试 warp/tests/test_external_build.py 验证了base.merged(addon)的完整结果例如 CUDA include 目录[cuda/base, shared, cuda/addon]shared只出现一次、preamble#define BASE_CUDA 1\n#define ADDON_CUDA 1\n且base本身保持不变。这使得独立开发的插件可以安全组合不会互相覆盖或复制全部选项。4.4 如何应用到模块set_module_optionsModuleBuildOptions通过既有模块选项 API 的新键extra_build_options应用该键在 warp/_src/context.py 被解析并归一化wp.set_module_options( {extra_build_options: wp.ModuleBuildOptions( extra_cuda_include_dirs[include_dir], extra_cuda_preamble#include my_addon.h, extra_build_dependencies[header_path], )}, modulesome_module, # 也接受 Module 对象或模块名 )set_module_options的module参数现在接受Module对象或模块名不再局限于 Python 模块。插件典型的使用模式design/external-compilation-extensions.md在 import 时注册语言元素类型与函数再把构建选项合并进每个消费模块的extra_build_options且必须在模块首次 JIT 启动或 AOT 编译之前完成。4.5 哈希与缓存失效规则模块哈希覆盖解析后的路径、preamble 文本、依赖文件内容、代码生成选项、以及被引用的 Warp 可见扩展契约外部 builtin 身份与原生类型码。两点关键行为Warp不会扫描传递的 C include——任何内容变化都应使缓存失效的头文件必须由插件显式列入extra_build_dependencies依赖内容在模块哈希被重新计算时wp.set_module_options()之后或新进程才被重读而不是每次 launch 都读——因此运行中的进程里修改头文件不会重建已加载模块。此外就地修改ModuleBuildOptions实例不会使已编译模块失效调用者必须通过wp.set_module_options()重新应用这是唯一的显式失效点warp/_src/context.py 与 design/external-compilation-extensions.md。测试 warp/tests/test_external_build.py 验证了依赖文件内容从#define EXTERNAL_VALUE 1改为2后模块哈希确实不同。5. 完整串联示例一个最小外部插件综合第 2~4 节一个最小插件的完整形态如下参考 warp/tests/test_external_build.pyimport ctypes import warp as wp # 1) 定义宿主 ABI 描述 class Color(ctypes.Structure): _fields_ [(r, ctypes.c_float), (g, ctypes.c_float), (b, ctypes.c_float)] # 2) import 时注册语言元素 wp.build_experimental.add_native_type( Color, native_namemy_addon::Color, fields{r: wp.float32, g: wp.float32, b: wp.float32}, initializeraggregate, ) wp.build_experimental.add_builtin( my_addon_scale, {value: Color, factor: wp.float32}, Color, native_namemy_addon::scale, ) # 3) 构造并应用模块级构建选项 PREAMBLE namespace my_addon { struct Color { float r, g, b; }; CUDA_CALLABLE inline Color scale(Color value, float factor) { return Color{value.r * factor, value.g * factor, value.b * factor}; } } wp.set_module_options( { extra_build_options: wp.ModuleBuildOptions( extra_cpu_preamblePREAMBLE, extra_cuda_preamblePREAMBLE, ), }, modulemy_app_module, # 消费该插件的模块 ) # 4) 在 kernel 中使用 wp.kernel(modulemy_app_module) def scale_kernel(colors: wp.array[Color], factor: wp.float32): i wp.tid() colors[i] wp.my_addon_scale(colors[i], factor)启动后即可wp.launch(scale_kernel, ...)正常执行。测试中还展示了原生值作为 kernel 参数时wp.Tape的反向传播不受影响不透明值作为普通参数参与梯度计算见test_native_value_tape_adjointwarp/tests/test_external_build.py。6. 扩展 ABI 与 AOT 产物路径warp.build_experimental之外外部编译特性还涉及两处配套 API 变化详见 design/external-compilation-extensions.md6.1 可选入口点 ABIexternal_constant_paramswp.kernel(entry_point_abi...)新增可选参数默认warp行为完全不变另一个取值是external_constant_params。该 ABI 生成一个无参数的extern CCUDA 入口点其唯一的 Warp struct 参数从模块级常量内存符号params中读取extern C { __constant__ __align__(alignof(MyParams)) unsigned char params[sizeof(MyParams)]; } extern C __global__ void my_kernel() { MyParams var_p *reinterpret_castconst MyParams*(params); // ... kernel body ... }入口点使用 kernel 的 mangled 名不带_cuda_kernel_forward后缀配合strip_hashTrue时Warp 使用去掉哈希后缀的 kernel key给外部运行时一个稳定的设备侧符号供查找。该 ABI 被刻意约束使用它的 kernel 必须仅限 CUDA、恰好接受一个 Warp struct 参数、设置enable_backwardFalse、避免wp.tid()、共享内存 tile 与确定性原子操作且不能传给wp.launch()。由于params是单一模块级符号同一模块内所有external_constant_paramskernel必须使用同一个 struct 类型混用类型会在代码生成期报错。6.2compile_aot_module返回产物路径wp.compile_aot_module()从返回None改为按目标顺序先设备、后显式arch值返回pathlib.Path对象列表。消费者不再需要重建 Warp 私有的缓存与架构命名规则忽略旧None返回值的调用者不受影响。7. 设计取舍与替代方案设计文档 design/external-compilation-extensions.md 记录了四个被否决的替代方案理解它们有助于把握 API 边界直接暴露内部注册 API代码改动小但 API 承诺大——内部add_builtin的大多数参数只服务于 Warp 自身的内置库没有稳定的外部契约进程级全局编译器设置全局add_header()/add_include_directory()/ 宏 / C 标准设置原型方便但后续每个模块都会继承变更插件无法隔离冲突设置缓存失效语义含糊并发编译还会观察到无关状态。模块级选项保留了实用能力且使所有权与哈希显式化注册源文件或库需要跨越 NVRTC、嵌入式 Clang、NVJitLink、目标架构、可重定位设备代码与缓存产物的公开链接模型头文件可见的实现已覆盖首批消费者独立源码/库支持可后续再加而不破坏契约把原生类型建模为 Warp struct 或新包装对象Warp struct 是外部定义的复制品、无法证明其 ABI新公开包装类型会在注解、数组、常量与 NumPy 互操作上引发身份与转换问题。ctypes 类本身已提供宿主布局无需额外对用户可见的类型。8. 限制与未来工作当前实验性 API 面被限定为模块构建输入、外部注册、入口点 ABI、AOT 产物路径返回契约。已知限制包括design/external-compilation-extensions.md注册表是进程全局的没有 unregister 机制动态注册的内置函数名对静态类型检查器不可见模块 reload 后新建的 ctypes 类与原始类在重载解析/参数匹配上不可互换不扫描传递 C include依赖头文件必须显式列出。文档列出的可行后续方向链接翻译单元、更强的 CUDA ABI 检查、外部 builtin 的梯度契约、更多构造策略、更多命名入口点 ABI、显式注册表生命周期——这些都应当跟随真实消费者出现而不是提前暴露编译器内部。9. 如何验证与测试仓库内的测试覆盖了本节 API 的全部行为可作为实现契约的权威参照warp/tests/test_external_build.py注册校验、冲突与幂等重注册、merged()顺序与非破坏性、依赖文件与外部 builtin 契约哈希、双后端 preamble 位置以及不透明/字段描述原生值作为 kernel 参数、builtin 结果、struct 字段与数组 dtype 的 NumPy 往返warp/tests/test_codegen.pypreamble 与后端专属 include 目录哈希、external_constant_params代码生成及其校验错误warp/tests/test_compilation.py 与 warp/tests/aot/test_module_aot.pyinclude 目录到达 Clang、CPU/CUDA 设备与架构列表的 AOT 产物路径warp/tests/cuda/test_occupancy.py 与 warp/tests/cuda/test_kernel_attributes.py外部入口点的占用率与 kernel 属性查询warp/tests/test_modules_lite.py确认内部add_builtin不属于公开 API。运行方式与仓库其他测试一致例如在仓库根目录执行python -m warp.tests.test_external_build或经由warp/tests/__main__.py的测试套件入口。设计文档同时提到针对真实外部消费者otk-pyoptix不 import 任何warp._src的验证目前为手动方式。最后提醒wp.build_experimental的experimental前缀是刻意的它让临时状态在 import 处可见。如果这些 API 未来达到 Warp 的稳定性与支持标准可能像warp.jax_experimental被提升为warp.jax一样被提升到稳定的warp.build命名空间——任何提升都会制定自己的兼容性与弃用计划。【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考