PyTorch C++ 稳定 ABI 算子注册宏完全指南：STABLE_TORCH_LIBRARY、STABLE_TORCH_LIBRARY_IMPL 与 TORCH_BOX

PyTorch C 稳定 ABI 算子注册宏完全指南STABLE_TORCH_LIBRARY、STABLE_TORCH_LIBRARY_IMPL 与 TORCH_BOX【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文以 PyTorch 稳定 ABIStable ABI的 C 算子注册宏为主题系统讲解STABLE_TORCH_LIBRARY、STABLE_TORCH_LIBRARY_IMPL、STABLE_TORCH_LIBRARY_FRAGMENT三个注册宏与TORCH_BOX包装宏的用途、参数和协作关系并结合 torch/csrc/stable/library.h 中的宏展开实现与 test/cpp_extensions/libtorch_agn_2_10_extension 中的真实测试代码讲清定义算子 → 注册内核的完整二进制兼容注册流程。读完后你可以独立编写一个在跨 PyTorch 版本间保持二进制兼容的 C 自定义算子扩展。一、为什么需要稳定 ABI 版注册宏标准的TORCH_LIBRARY、TORCH_LIBRARY_IMPL等宏依赖 PyTorch 内部 C 接口其符号与类型布局会随版本演进而变化。如果你的 C 扩展尤其是移动端、跨平台或第三方插件场景希望用旧版 PyTorch 的头文件编译链接到任意新版本的 libtorch 上运行就必须改用稳定 ABI 等价宏稳定 ABI 的 C 接口shim 函数是版本化、向前兼容的算子的 schema 定义与内核注册都通过稳定的 C 函数句柄完成而非直接调用内部 C 类。正如 注册宏文档所述这些宏是标准 PyTorch 算子注册宏TORCH_LIBRARY、TORCH_LIBRARY_IMPL等的稳定 ABI 等价物在需要跨 PyTorch 版本维持二进制兼容性的自定义算子中必须使用。四个宏STABLE_TORCH_LIBRARY、STABLE_TORCH_LIBRARY_IMPL、STABLE_TORCH_LIBRARY_FRAGMENT、TORCH_BOX的最低兼容版本均为 PyTorch 2.9。二、STABLE_TORCH_LIBRARY在命名空间中定义算子 schemaSTABLE_TORCH_LIBRARY(mylib, m) { m.def(my_op(Tensor input, int size) - Tensor); m.def(another_op(Tensor a, Tensor b) - Tensor); }参数说明参数含义ns算子命名空间如mylib即最终算子全名为mylib::my_opm在代码块内可用的StableLibrary变量名通过m.def(...)登记算子 schema关键约束每个命名空间只允许存在一个STABLE_TORCH_LIBRARY块。这一定位首次创建的语义在源码中体现为 StableLibrary 的构造函数会按Kind分派到不同的 C 接口——DEF走aoti_torch_library_init_def创建FRAGMENT走aoti_torch_library_init_fragment扩展IMPL走aoti_torch_library_init_impl注册内核。从源码结构看宏展开library.h#L355-L366生成一个静态初始化对象StableTorchLibraryInit其构造时立即调用你提供的 lambda 体fn(lib_)见 library.h#L130-L146随后StableLibrary析构时通过aoti_torch_delete_library_object释放 C 侧句柄。也就是说注册发生在进程加载静态初始化阶段而非某个函数被显式调用时。m.def()本身对应受限稳定版的torch::library::def()最终调用 C shimaoti_torch_library_deflibrary.h#L97-L101。在较新的目标版本下StableLibrary还提供了带 tag 的def重载与set_python_module能力分别由TORCH_FEATURE_VERSION的 2.12 / 2.13 版本门控见 library.h#L103-L127说明该 API 面本身也在按稳定 ABI 的纪律渐进扩展。三、STABLE_TORCH_LIBRARY_IMPL为 dispatch key 注册内核STABLE_TORCH_LIBRARY_IMPL(mylib, CPU, m) { m.impl(my_op, TORCH_BOX(my_cpu_kernel)); } STABLE_TORCH_LIBRARY_IMPL(mylib, CUDA, m) { m.impl(my_op, TORCH_BOX(my_cuda_kernel)); }参数说明参数含义ns算子所在命名空间须与STABLE_TORCH_LIBRARY一致k目标 dispatch key如CPU、CUDA、CompositeExplicitAutograd等m块内可用的StableLibrary变量通过m.impl(name, fn)绑定内核核心要点所有通过此宏注册的内核函数都必须先用TORCH_BOX包装。原因在于稳定 ABI 侧内核统一采用boxed calling convention——函数签名为void (*fn)(StableIValue*, uint64_t num_inputs, uint64_t num_outputs)输入输出都以StableIValue值栈传递而不是任意 C 类型参数。m.impl()的稳定实现见 library.h#L85-L95它校验 schema 声明的实参/输出数量与内核函数签名一致不匹配时抛出Registered schema has N args, but the kernel to box has M之类的错误见 library.h#L227-L248然后将其登记到 C 侧库对象中。宏展开细节library.h#L337-L353值得注意STABLE_TORCH_LIBRARY_IMPL(ns, k, m)内部借助C10_UID生成唯一符号名_STABLE_TORCH_LIBRARY_IMPL因此同一个翻译单元内可以针对同一命名空间注册多个不同 dispatch key而STABLE_TORCH_LIBRARY不带C10_UID正是每命名空间仅一次限制的实现来源。四、STABLE_TORCH_LIBRARY_FRAGMENT跨翻译单元扩展同一命名空间STABLE_TORCH_LIBRARY_FRAGMENT(mylib, m) { m.def(extra_op(Tensor x) - Tensor); }它是TORCH_LIBRARY_FRAGMENT的稳定 ABI 等价物用于向已经由STABLE_TORCH_LIBRARY创建的命名空间追加算子定义典型场景是多文件扩展mylib在第一个 .cpp 里用STABLE_TORCH_LIBRARY创建其余 .cpp 均用STABLE_TORCH_LIBRARY_FRAGMENT补充定义。其展开library.h#L368-L383同样基于C10_UID避免符号冲突并以Kind::FRAGMENT触发aoti_torch_library_init_fragment。五、TORCH_BOX把普通 C 函数适配到 boxed 调用约定Tensor my_kernel(const Tensor input, int64_t size) { return input.reshape({size}); } STABLE_TORCH_LIBRARY_IMPL(my_namespace, CPU, m) { m.impl(my_op, TORCH_BOX(my_kernel)); }TORCH_BOX(func)接收一个未装箱unboxed内核函数指针生成符合 boxed 调用约定的包装函数。从 library.h#L332-L335 的定义看它只是实例化模板torch::stable::detail::boxerFuncT, func::boxed_fn通过infer_function_traits_t推导返回类型与参数列表对多返回值std::tuple、单返回值、void 返回值三种情形分别给出特化library.h#L205-L311 实际位于 torch/csrc/stable/library.h#L205-L311运行时先校验num_args/num_outputs与函数签名匹配再用unbox_to_tuple把StableIValue栈转为 C 参数、调用原函数、最后用box_from_tuple把结果写回栈中类型层面UnboxType做了一些所有权语义适配例如HeaderOnlyArrayRefT会映射为std::vectorT、std::string_view映射为std::stringlibrary.h#L148-L176保证从值栈取出的数据在函数返回后依然有效。这解释了为什么你的内核函数必须使用稳定 ABI 头文件里的类型如torch::stable::Tensor而不是普通at::Tensor——toT/fromT转换只对稳定类型集开放。六、真实示例测试扩展中的完整注册写法仓库中 test/cpp_extensions/libtorch_agn_2_10_extension 下有一批面向目标 PyTorch 2.10 的旧 ABI 扩展的测试用例每个文件都是一份最小可复制的模板。以 my_reshape.cpp 为例#include torch/csrc/stable/library.h #include torch/csrc/stable/tensor.h #include torch/csrc/stable/ops.h using torch::stable::Tensor; Tensor my_reshape(Tensor t, torch::headeronly::HeaderOnlyArrayRefint64_t shape) { return reshape(t, shape); // 调用稳定 ops.h 提供的接口 } STABLE_TORCH_LIBRARY_FRAGMENT(STABLE_LIB_NAME, m) { m.def(my_reshape(Tensor t, int[] shape) - Tensor); } STABLE_TORCH_LIBRARY_IMPL(STABLE_LIB_NAME, CompositeExplicitAutograd, m) { m.impl(my_reshape, TORCH_BOX(my_reshape)); }几个实践要点头文件只 include 稳定 ABI 头torch/csrc/stable/library.h、torch/csrc/stable/tensor.h、torch/csrc/stable/ops.h等这是二进制兼容的前提可选参数用std::optional如 my_full.cpp 中std::optionaltorch::headeronly::ScalarType dtype对应 schema 里的ScalarType? dtypeNonedispatch key 的选择这些测试多注册到CompositeExplicitAutograd用已有算子组合实现、自动获得 autograd 支持而非CPU/CUDA这类底层设备 key该目录内my_bitwise.cpp、my_permute.cpp、my_full.cpp等文件可当作一个算子一个翻译单元、各自用 FRAGMENT IMPL 注册的批量参考。七、宏实现与版本门控的底层细节注册时机三个STABLE_*宏展开后都创建一个static const StableTorchLibraryInit对象library.h#L130-L146在动态库加载/静态初始化阶段执行注册 lambda因此无需在main中手动触发。C/C 边界StableLibrary内部只持有TorchLibraryHandle lib_一个不透明 C 句柄所有def/impl都降级为 C shim 调用aoti_torch_library_def、torch_library_impl等析构时aoti_torch_delete_library_object释放——这正是稳定 ABI的落地形态跨版本只承诺 C 接口稳定。版本门控StableLibrary::impl在TORCH_FEATURE_VERSION 2.10时改走带 ABI 版本号的torch_library_impl否则回退到aoti_torch_library_impllibrary.h#L88-L93。类似的版本自适应还体现在 torch/csrc/stable/macros.h 中TORCH_DYNAMIC_VERSION_CALL宏通过运行时符号查找dlsym/GetProcAddress在旧目标扩展上调用新 libtorch 才提供的 shim找不到则使用签名一致的 fallback 函数macros.h#L83-L126STABLE_TORCH_ERROR_CODE_CHECK则负责把 C 侧错误码转换为携带原始错误信息的 C 异常macros.h#L128-L168。适用前提与限制文档标注的最小兼容版本为 PyTorch 2.9本文所有源码分析基于当前仓库实现。稳定 ABI 头文件仅覆盖其支持的类型与算子子集编写扩展时应以 torch/csrc/stable 目录下的头文件library.h、tensor.h、ops.h、device.h、generator.h等为准避免引入普通 libtorch 内部类型。八、速查表宏作用关键约束STABLE_TORCH_LIBRARY(ns, m)创建命名空间并定义算子 schema每命名空间仅一个m.def(schema_string)STABLE_TORCH_LIBRARY_FRAGMENT(ns, m)向已有命名空间追加算子定义命名空间须已存在跨翻译单元扩展用STABLE_TORCH_LIBRARY_IMPL(ns, k, m)为 dispatch keyk注册内核内核必须经TORCH_BOX装箱同一命名空间可注册多个 keyTORCH_BOX(func)将普通 C 内核函数适配为 boxed 调用约定参数/返回类型须为稳定 ABI 支持的类型以上四个宏最低兼容版本均为 PyTorch 2.9配合 torch/csrc/stable 下的稳定头文件即可构建跨 PyTorch 版本二进制兼容的 C 自定义算子扩展。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考