PyTorch Stable C API 完全指南二进制兼容的算子注册与跨版本 AOT 扩展开发【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文基于 PyTorch 仓库中的官方 C 文档docs/cpp/source/api/stable/展开系统讲解 Stable C API 的设计理念、注册宏用法、稳定算子集合与工具类能力并结合torch/csrc/stable/与torch/headeronly/的源码实现说明 boxed kernel 调用约定、版本定向version targeting等底层机制帮助你编写一次编译、跨 PyTorch 版本运行的 C 扩展。一、Stable C API 是什么PyTorch Stable C API 提供了一组二进制兼容的接口用于调用张量操作与常用工具并保证这些接口在 PyTorch 各版本之间保持稳定。其核心价值是使用它编译出的 ahead-of-timeAOT扩展在升级 PyTorch 时无需重新编译即可继续运行。官方文档明确给出了四类典型适用场景构建需要在多个 PyTorch 版本上工作的扩展分发预编译二进制二进制兼容性比第一时间获得最新特性更重要为生产环境部署编写自定义算子。最小基本用法如下继承自docs/cpp/source/api/stable/index.md#include torch/csrc/stable/library.h #include torch/csrc/stable/ops.h // Create a tensor using stable API auto tensor torch::stable::empty( {3, 4}, torch::headeronly::ScalarType::Float, torch::headeronly::Layout::Strided, torch::stable::Device(torch::headeronly::DeviceType::CPU), false, torch::headeronly::MemoryFormat::Contiguous); // Register operators with stable ABI STABLE_TORCH_LIBRARY(myops, m) { m.def(my_op(Tensor input) - Tensor); } STABLE_TORCH_LIBRARY_IMPL(myops, CPU, m) { m.impl(my_op, TORCH_BOX(my_cpu_kernel)); }注意这里的 API 命名空间划分稳定张量操作位于torch::stable::而标量类型、设备类型、内存格式等枚举与基础类型则来自torch::headeronly::命名空间——这是后文“Header-Only 工具”一节的关键。二、头文件组织官方文档列出的六个核心头文件在仓库中均真实存在于 torch/csrc/stable/ 目录头文件职责torch/csrc/stable/library.h稳定库注册StableLibrary类与注册宏torch/csrc/stable/ops.h稳定算子定义torch::stable::*函数集torch/csrc/stable/tensor.h稳定张量结构torch::stable::Tensortorch/csrc/stable/device.h稳定设备结构torch::stable::Devicetorch/csrc/stable/accelerator.h加速器支持DeviceGuard、Stream等torch/csrc/stable/macros.h稳定 API 宏CUDA 错误检查等从 torch/csrc/stable/library.h 的头部注释可以看出设计取向该文件“只能包含稳定内容”并且与底层 C shim 不同它允许包含 header-only C 代码以提升用户体验——即 Stable API 是构建在torch/csrc/stable/c/shim.h与torch/csrc/inductor/aoti_torch/c/shim.h这些 ABI 稳定的 C 接口之上的 C 封装层。三、库注册宏STABLE_TORCH_LIBRARY 家族注册宏提供标准 PyTorch 算子注册宏TORCH_LIBRARY、TORCH_LIBRARY_IMPL等的 stable ABI 等价物用于构建需要在 PyTorch 版本间保持二进制兼容的自定义算子。3.1 STABLE_TORCH_LIBRARY定义算子 schemaSTABLE_TORCH_LIBRARY(ns, m)是TORCH_LIBRARY的稳定等价宏用于在指定命名空间中定义算子 schema。约束是每个命名空间只能有一个STABLE_TORCH_LIBRARY块若需在不同翻译单元中为同一命名空间追加定义应使用STABLE_TORCH_LIBRARY_FRAGMENT。STABLE_TORCH_LIBRARY(mylib, m) { m.def(my_op(Tensor input, int size) - Tensor); m.def(another_op(Tensor a, Tensor b) - Tensor); }最小兼容版本PyTorch 2.9。3.2 STABLE_TORCH_LIBRARY_IMPL注册 dispatch key 实现STABLE_TORCH_LIBRARY_IMPL(ns, k, m)为特定 dispatch key如CPU、CUDA注册算子实现。关键约束是所有通过该宏注册的 kernel 函数都必须用TORCH_BOX宏装箱。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)); }最小兼容版本PyTorch 2.9。3.3 STABLE_TORCH_LIBRARY_FRAGMENT 与 TORCH_BOXSTABLE_TORCH_LIBRARY_FRAGMENT(ns, m)是TORCH_LIBRARY_FRAGMENT的稳定等价物用于向已由STABLE_TORCH_LIBRARY创建的命名空间追加算子定义PyTorch 2.9。TORCH_BOX(func)将普通unboxedkernel 函数指针包装为符合 stable boxed kernel 调用约定的 wrapperTensor 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/csrc/stable/library.h 中TORCH_BOX被实现为对元编程模板boxer的取址#define TORCH_BOX(func) \ torch::stable::detail::boxer \ std::remove_pointer_tstd::remove_reference_tdecltype(func), \ (func)::boxed_fn其内部的boxer_impl见 library.h按返回值分为三种特化多返回std::tuple...、单返回、void。装箱后的函数签名为void(StableIValue* stack, uint64_t num_args, uint64_t num_outputs)运行时先校验 schema 声明的参数/输出个数与 kernel 实际签名是否一致不一致会抛错再把StableIValue栈按 schema 解包为真实 C 类型unbox_to_tuple、调用用户函数、把结果重新打包回栈box_from_tuple。此外UnboxType模板library.h会特地把HeaderOnlyArrayRefT映射为std::vectorT、string_view映射为string以保留所有权语义。这正是稳定 ABI 中“boxed kernel 调用约定”的具体含义。在StableLibrary类层面library.hdef()与impl()最终都落到 C shim 上aoti_torch_library_def与torch_library_impl/aoti_torch_library_impl注册失败时通过STABLE_TORCH_ERROR_CODE_CHECK宏抛出携带原始错误信息的std::runtime_error。另外StableLibrary::Kind枚举区分DEF/IMPL/FRAGMENT三种库初始化路径分别调用aoti_torch_library_init_def、aoti_torch_library_init_impl、aoti_torch_library_init_fragment。源码中还可见两个版本演进点文档未展开此处以源码为准PyTorch 2.12 提供带 tag 的def(schema, tags)重载底层调用torch_library_def_with_tagsPyTorch 2.13 提供set_python_module(pymodule, context)用于把 Python 模块与稳定库关联。四、Stable 算子集合Stable API 提供的张量操作均保持跨版本二进制兼容。按文档docs/cpp/source/api/stable/operators.md的分类完整清单如下且均可在 torch/csrc/stable/ops.h 中以inline函数形式找到对应实现Tensor 与 Device 类torch::stable::Tensor tensor torch::stable::empty({3, 4}, ...); float* data tensor.data_ptrfloat(); auto shape tensor.sizes(); torch::stable::Device cpu_device(torch::headeronly::DeviceType::CPU); torch::stable::Device cuda_device(torch::headeronly::DeviceType::CUDA, 0);张量创建torch::stable::empty、torch::stable::empty_liketorch::stable::new_empty(self, size, dtype?, layout?, device?, pin_memory?)torch::stable::new_zeros(self, size, dtype?, layout?, device?, pin_memory?)torch::stable::fulltorch::stable::from_blob(data, sizes, strides, device, dtype, storage_offset, layout)创建示例auto tensor torch::stable::empty( {3, 4}, torch::headeronly::ScalarType::Float, torch::headeronly::Layout::Strided, torch::stable::Device(torch::headeronly::DeviceType::CUDA, 0), false, torch::headeronly::MemoryFormat::Contiguous);张量变换clone、contiguous、reshape、view按 size 与按 dtype 两个重载、permute、flatten、squeeze、unsqueeze、transpose、select、index_select、narrow、pad。设备与类型转换to全参数重载与 Device 重载两个版本、is_pinned。原地操作fill_、zero_、copy_。数学运算matmul、amax单维与多维两个重载、sum、sum_out、subtract、bitwise_and、bitwise_or、bitwise_left_shift、bitwise_right_shift、floor_divide。从 torch/csrc/stable/ops.h 的实现结构看这些函数都是对底层 stable shim 的薄封装例如fill_、narrow、amax、transpose、matmul均直接以inline函数提供参数校验依赖 header-only 的STD_TORCH_CHECK从而避免了对 libtorch 符号的链接依赖。五、工具类设备守护、流与 CUDA 错误检查这一节完整继承 docs/cpp/source/api/stable/utilities.md 的内容。5.1 DeviceGuard 与当前设备torch::stable::accelerator::DeviceGuard是 RAII 设备切换器getCurrentDeviceIndex获取当前设备索引{ torch::stable::accelerator::DeviceGuard guard(1); // Operations here run on device 1 } // Previous device is restored5.2 获取当前 CUDA Stream分版本写法PyTorch 2.13使用accelerator.h中的高层 API#include torch/csrc/stable/accelerator.h // nativeHandle() requires PyTorch 2.13 cudaStream_t stream static_castcudaStream_t( torch::stable::accelerator::getCurrentStream(tensor.get_device_index()).nativeHandle()); // Now you can use stream in your CUDA kernel launches my_kernelblocks, threads, 0, stream(args...);PyTorch 2.9–2.12通过 ABI 稳定的 C shim API 获取#include torch/csrc/inductor/aoti_torch/c/shim.h #include torch/headeronly/util/shim_utils.h // Use the ABI-stable C shim API to get the current CUDA stream. void* stream_ptr nullptr; TORCH_ERROR_CODE_CHECK( aoti_torch_get_current_cuda_stream(tensor.get_device_index(), stream_ptr)); cudaStream_t stream static_castcudaStream_t(stream_ptr); // Now you can use stream in your CUDA kernel launches my_kernelblocks, threads, 0, stream(args...);官方特别强调使用 C shim API 时必须用TORCH_ERROR_CODE_CHECK宏检查错误码并抛出异常而nativeHandle()这类高层 C 工具 API 已替你做了这一检查。5.3 CUDA 错误检查宏PyTorch 2.10STD_CUDA_CHECK(EXPR)检查 CUDA API 调用结果出错时抛异常使用前需自行包含cuda_runtime.hSTD_CUDA_KERNEL_LAUNCH_CHECK()检查最近一次 kernel 启动的错误等价于STD_CUDA_CHECK(cudaGetLastError())STD_CUDA_CHECK(cudaMalloc(ptr, size)); STD_CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost)); my_kernelblocks, threads, 0, stream(args...); STD_CUDA_KERNEL_LAUNCH_CHECK();源码实现位于 torch/csrc/stable/macros.hSTD_CUDA_CHECK通过 stable shimtorch_c10_cuda_check_msg生成使用 PyTorch 错误格式化的详细消息并以std::runtime_error抛出整个宏组被#if TORCH_FEATURE_VERSION TORCH_VERSION_2_10_0版本门控与文档标注的“最小兼容版本 PyTorch 2.10”一致。5.4 并行化工具torch::stable::parallel_for跨版本兼容的 CPU 并行循环入口torch::stable::get_num_threads获取线程数。两者在 torch/csrc/stable/ops.h 中均有对应 inline 实现约 L570 与 L594 附近。六、Header-Only 工具彻底摆脱 libtorch 链接依赖torch::headeronly命名空间提供常用 PyTorch 类型与工具的纯头文件版本完全无需链接 libtorch——这一可移植性正是维持跨版本二进制兼容的理想选择。torch/headeronly/README.md 说明这些头文件与 LibTorch 完全解耦且全部 API 全局列在 torch/header_only_apis.txt 中。6.1 错误检查STD_TORCH_CHECK#include torch/headeronly/util/Exception.h STD_TORCH_CHECK(condition, Error message with , variable, interpolation);凡原来使用TORCH_CHECK的位置都可以换成STD_TORCH_CHECK以消除 libtorch 链接。唯一差异条件不满足时TORCH_CHECK抛出更精致的c10::Error而STD_TORCH_CHECK抛出std::runtime_error。6.2 核心类型以下c10::类型在torch::headeronly::下有 header-only 版本torch::headeronly::ScalarType——张量数据类型Float、Double、Int 等torch::headeronly::DeviceType——设备类型CPU、CUDA 等torch::headeronly::MemoryFormat——内存布局Contiguous、ChannelsLast 等torch::headeronly::Layout——张量布局Strided、Sparse 等#include torch/headeronly/core/ScalarType.h #include torch/headeronly/core/DeviceType.h #include torch/headeronly/core/MemoryFormat.h #include torch/headeronly/core/Layout.h auto dtype torch::headeronly::ScalarType::Float; auto device_type torch::headeronly::DeviceType::CUDA; auto memory_format torch::headeronly::MemoryFormat::Contiguous; auto layout torch::headeronly::Layout::Strided;6.3 TensorAccessor从稳定张量的数据指针、尺寸与步长构造带边界检查的高效访问器#include torch/headeronly/core/TensorAccessor.h // Create a TensorAccessor for a 2D float tensor auto sizes tensor.sizes(); auto strides tensor.strides(); torch::headeronly::TensorAccessorfloat, 2 accessor( static_castfloat*(tensor.mutable_data_ptr()), sizes.data(), strides.data()); // Access elements float value accessor[i][j];6.4 Dispatch 宏THO Torch Header OnlyTHO_DISPATCH_V2与AT_DISPATCH_V2见ATen/Dispatch_v2.h行为一致但不需要链接 libtorch对未实现的 dtype 路径AT_DISPATCH_V2抛c10::NotImplementedErrorTHO_DISPATCH_V2抛std::runtime_error。#include torch/headeronly/core/Dispatch_v2.h THO_DISPATCH_V2( tensor.scalar_type(), // will be resolved as scalar_t my_kernel, AT_WRAP(([]() { // code to specialize with scalar_t // scalar_t is the resolved C type (e.g. float, double) auto* data static_castscalar_t*(tensor.mutable_data_ptr()); Scalar s(*data); })), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_COMPLEX_TYPES), torch::headeronly::ScalarType::Half, // as many type arguments as needed );同时一批 AT_* 类型集合宏已迁移为 header-onlyAT_FLOATING_TYPES、AT_INTEGRAL_TYPES、AT_INTEGRAL_TYPES_V2、AT_ALL_TYPES、AT_COMPLEX_TYPES、AT_ALL_TYPES_AND_COMPLEX、AT_FLOAT8_TYPES、AT_BAREBONES_UNSIGNED_TYPES、AT_QINT_TYPES。对于仍使用旧版v1AT_DISPATCH体系的扩展迁移是机械式的AT_DISPATCH_SWITCH→THO_DISPATCH_SWITCH、AT_DISPATCH_CASE→THO_DISPATCH_CASE、AT_PRIVATE_CASE_TYPE_USING_HINT→THO_PRIVATE_CASE_TYPE_USING_HINT、at::ScalarType::X→torch::headeronly::ScalarType::X。文档给出的完整迁移前后对照示例// ---- Before (requires linking against libtorch) ---- #include torch/all.h #define MY_DISPATCH_CASE_FLOATING_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) #define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ AT_DISPATCH_SWITCH(TYPE, NAME, \ MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))// ---- After (header-only, no libtorch dependency) ---- #include torch/headeronly/core/Dispatch.h #define MY_DISPATCH_CASE_FLOATING_TYPES(...) \ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) #define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ THO_DISPATCH_SWITCH(TYPE, NAME, \ MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))七、版本定向TORCH_TARGET_VERSION一个文档未覆盖但源码中明确存在的重要机制是稳定 ABI 版本定向定义在 torch/csrc/stable/version.h默认情况下扩展使用编译期 libtorch 头文件的当前TORCH_ABI_VERSION也可以显式指定目标版本使同一份代码在不同 libtorch 头文件下编译出面向旧 ABI 的扩展// (1) 编译器参数-DTORCH_TARGET_VERSION0x0209000000000000 // (2) 或在 include 任何头文件之前于源码中定义 #define TORCH_TARGET_VERSION (((0ULL 2) 56) | ((0ULL 9) 48)) #include torch/csrc/stable/library.h硬约束TORCH_TARGET_VERSION必须小于等于所用 libtorch 头文件的TORCH_ABI_VERSION否则直接编译报错。头文件内还定义了TORCH_VERSION_2_10_0至TORCH_VERSION_2_14_0等版本常量源码中各 API 以#if TORCH_FEATURE_VERSION TORCH_VERSION_x_y_0进行版本门控——这解释了前文“2.10 才有 STD_CUDA_CHECK”“2.12 才有 tag 注册”“2.13 才有 nativeHandle/set_python_module”等版本要求的出处。此外 torch/csrc/stable/macros.h 提供了TORCH_DYNAMIC_VERSION_CALL机制对目标版本较旧的扩展可在运行时用dlsym/GetProcAddress探测当前 libtorch 是否已提供新增 shim若有则动态调用、否则回退到 fallback 函数。官方注释强调该宏开销较大只应在“旧目标扩展确实需要新 shim”的罕见场景使用。八、验证与非稳定 API 的对照仓库中的测试覆盖了 stable ABI 的 iValue 转换与扩展构建test/cpp/shim/test_stable_ivalue.cpp、test/cpp_extensions/libtorch_agn_2_10_extension/含跨版本扩展示例 my_stable_error_check.cpp。若你的场景不要求二进制兼容应改用标准的非稳定算子注册 API参见 docs/cpp/source/api/library/index.md。小结Stable C API 的选型逻辑很清晰——用STABLE_TORCH_LIBRARY/STABLE_TORCH_LIBRARY_IMPLTORCH_BOX注册算子用torch::stable::*完成张量操作用torch::headeronly::*完成无需链接 libtorch 的类型、检查与 dispatch并借助TORCH_TARGET_VERSION锁定 ABI 目标版本即可构建一次编译、多版本运行的生产级 C 扩展。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考