第7章 PyTorch 自定义算子指南

第7章 PyTorch 自定义算子 创建时间：2024年6月18日 | 最后更新：2026年6月16日 | 最后验证：2024年11月05日

PyTorch 提供了大量可用于 Tensor 的内置算子（例如 `torch.add`、`torch.sum` 等）。不过，若你想引入新的自定义操作，并使其与 `torch.compile`、`autograd` 或 `torch.vmap` 等子系统协同工作，则必须通过 Python 的 `torch.library` 文档或 C++ 的 `TORCH_LIBRARY` API 将该操作注册到 PyTorch。

从 Python 编写自定义算子

请参阅 Python 自定义算子指南。该指南涵盖了必需的 schema 及 mutation/aliasing 契约、opcheck、函数式算子、可变算子，以及 `autograd` 和 `torch.vmap` 等可选注册项。

若符合以下情况，建议从 Python 而非 C++ 编写自定义算子： - 你希望 PyTorch 将某个 Python 函数视为不透明的可调用对象，特别是在涉及 `torch.compile` 和 `torch.export` 时。 - 你拥有 C++/CUDA kernel 的 Python 绑定，希望这些绑定能与 PyTorch 子系统（如 `torch.compile` 或 `torch.autograd`）组合使用。 - 你使用的是 Python 环境（而非 AOTInductor 等纯 C++ 环境）。

将自定义 C++ 和/或 CUDA 代码集成到 PyTorch

请参阅 C++ 和 CUDA 自定义算子指南。

注意 SYCL 是 Intel GPU 的后端编程语言。有关自定义 SYCL 代码的集成，请参阅 SYCL 自定义算子指南。

若符合以下情况，建议从 C++ 而非 Python 编写自定义算子： - 你拥有自定义的 C++ 和/或 CUDA 代码。 - 你计划使用 AOTInductor 进行无 Python 依赖的推理。

自定义算子手册

对于教程和本页面未涵盖的信息，请参阅自定义算子手册。我们建议先阅读上述任一教程，再将手册作为参考查阅；手册并非按顺序通读的材料。

何时应创建自定义算子？

如果你的操作可以由内置 PyTorch 算子组合而成，请将其编写为 Python 函数并直接调用，而非创建自定义算子。仅在调用 PyTorch 无法理解的库（例如自定义 C/C++ 代码、自定义 CUDA kernel 或 C/C++/CUDA 扩展的 Python 绑定）时，才使用算子注册 API 创建自定义算子。

为何要创建自定义算子？

虽然可以通过获取 Tensor 的数据指针并将其传递给 pybind 封装的 kernel 来调用 C/C++/CUDA kernel，但这种方法无法与 `autograd`、`torch.compile`、`vmap` 等 PyTorch 子系统良好组合。为了让操作能与 PyTorch 各子系统组合，必须通过算子注册 API 进行注册。