PyTorch 自定义操作深度解读：从 C++/CUDA 实现到 AOTInductor 部署

背景

在 PyTorch 模型的开发与部署过程中，使用自定义操作（Custom Operations）已成为一种常见且必要的手段。无论是为了追求极致的推理性能，还是为了实现特定硬件上的优化，开发者往往需要绕过 PyTorch 内置算子，直接利用 C++ 和 CUDA 编写底层逻辑。

然而，将自定义操作无缝集成到 PyTorch 生态系统中并非易事。这不仅涉及在 C++ 和 CUDA 层面实现自定义类（Custom Classes）和自定义函数（Custom Functions），还要求这些组件能够同时兼容 Python 和 C++ 的推理程序。此外，随着 torch.compile 和 torch.export 等编译技术的普及，如何确保自定义操作在符号追踪（Symbolic Tracing）和 Ahead-of-Time (AOT) 编译过程中正确工作，成为了一个关键的技术挑战。

本文基于 Hacker News 上关于 PyTorch Custom Operation 的技术讨论，通过一个简单的“恒等卷积”（Identity Convolution）示例，深入解析如何在 C++ 和 CUDA 中实现 PyTorch 自定义操作，并展示其在 PyTorch 模型及 AOTInductor 编译后的推理程序中的完整使用流程。

核心内容

1. PyTorch 自定义函数（Custom Function）

PyTorch 自定义函数通常用于实现无状态的计算逻辑。这类函数可以使用 C++ 和 CUDA 进行实现，并通过 TORCH_LIBRARY_IMPL 宏进行注册。

• 多设备支持：开发者可以同时提供 CPU 和 CUDA 两种实现。PyTorch 会根据输入张量（Tensor）所在的设备（Device），自动将计算分发（Dispatch）到正确的实现版本上。
• 无状态特性：自定义函数本身不持有参数或状态，仅执行纯计算任务。

2. PyTorch 自定义类（Custom Class）

与自定义函数不同，自定义类用于实现有状态的操作，例如持有模型参数或内部状态。

• 实现方式：若需实现一个包含参数且具有 forward() 方法的自定义类，使其能从 Python 端调用，可以使用 torch::CustomClassHolder 在 C++ 中定义该类，并通过 TORCH_LIBRARY 宏进行注册。
• 应用场景：适用于需要维护内部状态或复杂对象生命周期的场景。

3. 在 PyTorch 中使用自定义操作和类

自定义类和函数的 C++ 实现会被编译成一个共享库（例如 libidentity_conv_ops.so）。要在 PyTorch 中加载和使用这些组件，需遵循以下步骤：

• 加载库：使用 torch.ops.load_library 加载共享库。加载后，自定义类可通过 torch.classes 访问，自定义函数可通过 torch.ops 访问。
• 兼容 torch.compile 和 torch.export：
  • 为了确保基于 FakeTensor 的符号追踪能正常工作，必须在 PyTorch 中注册自定义类和函数的“虚假”（Fake，即抽象）版本。
  • 这可以通过 @register_fake_class 和 @torch.library.register_fake 装饰器/宏来实现。
  • 目的：在追踪阶段，PyTorch 无需执行实际的 C++/CUDA 代码，而是使用这些抽象版本进行符号表示，从而避免在追踪过程中触发实际的硬件计算或依赖问题。

4. PyTorch 模型导出与降级（Export and Lowering）

当模型中使用了自定义类和自定义函数时，其导出和编译流程如下：

• 模型导出：
  • 如果已为所有自定义类和函数注册了用于 torch.export 符号追踪的抽象版本，则可以使用 torch.export 导出模型。
  • 图表示：在导出的计算图中，自定义类 IdentityConvClass.forward 会被表示为对 torch.ops.higher_order.call_torchbind 的调用；而自定义操作 identity_conv_op 则被表示为对 torch.ops.my_ops.identity_conv_op 的调用。
• 编译与打包：
  • 使用 torch._inductor.aoti_compile_and_package 对导出后的程序进行编译和打包，生成 model.pt2 包。
  • 该包可以被 Python 和 C++ 推理程序加载。
• 运行时行为：
  • 当编译后的模型执行时，自定义类和自定义操作的实现会从共享库中加载，并在运行时被正确分发（Dispatch）。
• 纯 C++ 推理支持：
  • 在纯 C++ 推理程序中，可以通过 dlopen 加载自定义类和函数的共享库并进行注册。
  • 优势：此过程不需要 pybind11 或 libpython 依赖，实现了与 Python 环境的解耦，有利于部署在资源受限或纯 C++ 环境中。

关键要点

• 实现与注册分离：C++/CUDA 实现需编译为共享库，并通过 TORCH_LIBRARY_IMPL（函数）和 TORCH_LIBRARY（类）进行注册。
• 设备自动分发：PyTorch 能根据输入张量设备自动选择 CPU 或 CUDA 实现，无需手动干预。
• 符号追踪的关键：必须注册自定义组件的“Fake”（抽象）版本，以支持 torch.export 和 torch.compile 的符号追踪，避免追踪阶段执行实际代码。
• 图结构变化：导出后，自定义类调用映射为 call_torchbind，自定义操作映射为具体的 torch.ops 调用。
• AOT 编译优势：通过 aoti_compile_and_package 生成的 .pt2 包支持 Python 和 C++ 双端推理。
• 无 Python 依赖部署：纯 C++ 环境可通过 dlopen 加载共享库，无需链接 libpython 或 pybind11，降低了部署复杂度。

意义与影响

这篇技术解读揭示了 PyTorch 在高性能计算和灵活部署方面的核心机制。对于 AI 工程师和系统开发者而言，掌握自定义操作（Custom Ops）的实现与集成，意味着能够突破框架内置算子的性能瓶颈，针对特定硬件或算法需求进行极致优化。

特别是关于 torch.export 和 AOTInductor 的讨论，强调了在现代 PyTorch 工作流中，**“可编译性”**的重要性。通过引入 FakeTensor 机制和抽象注册，PyTorch 成功地将动态的 Python 执行环境与静态的 C++ 编译优化连接起来，使得自定义代码也能享受到 torch.compile 带来的性能红利。

此外，支持纯 C++ 推理且无 Python 依赖的特性，极大地扩展了 PyTorch 模型的部署边界。这使得 PyTorch 模型不仅能在服务器端通过 Python 服务运行，也能高效地嵌入到游戏引擎、嵌入式设备或高性能 C++ 推理引擎中，为 AI 技术的广泛落地提供了坚实的技术基础。