PyTorch 自定义算子#

创建日期：2024 年 6 月 18 日 | 最后更新：2026 年 6 月 16 日 | 最后验证：2024 年 11 月 5 日

PyTorch 提供了一个庞大的算子库来处理张量（例如 torch.add、torch.sum 等）。然而，您可能希望将新的自定义操作引入 PyTorch，并使其与 torch.compile、自动微分（autograd）和 torch.vmap 等子系统配合使用。为此，您必须通过 Python torch.library 文档或 C++ TORCH_LIBRARY API 将自定义操作注册到 PyTorch 中。

使用 Python 编写自定义算子#

请参阅自定义 Python 算子获取 Python 指南。它涵盖了所需的模式（schema）、变异/别名约定、opcheck、函数式算子、可变算子，以及自动微分和 torch.vmap 等可选注册项。

在以下情况下，您可能希望使用 Python（而不是 C++）来编写自定义算子：

您有一个希望 PyTorch 将其视为不透明可调用对象的 Python 函数，特别是在涉及 torch.compile 和 torch.export 时。
您拥有一些指向 C++/CUDA 内核的 Python 绑定，并希望将其与 PyTorch 子系统（如 torch.compile 或 torch.autograd）进行组合。
您正在使用 Python（而不是像 AOTInductor 那样的纯 C++ 环境）。

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

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

注意

SYCL 是英特尔 GPU 的后端编程语言。如需集成自定义 Sycl 代码，请参阅自定义 SYCL 算子。

在以下情况下，您可能希望使用 C++（而不是 Python）来编写自定义算子：

您拥有自定义 C++ 和/或 CUDA 代码。
您计划将此代码与 AOTInductor 一起使用，以进行脱离 Python 的推理。

自定义算子手册#

对于教程和本页面未涵盖的信息，请参阅自定义算子手册。我们建议您先阅读上述教程之一，然后将自定义算子手册作为参考；它并不适合从头到尾阅读。

什么时候应该创建自定义算子？#

如果您的操作可以通过内置 PyTorch 算子的组合来表示，请将其编写为 Python 函数并调用，而不是创建自定义算子。只有在调用 PyTorch 无法理解的库（例如自定义 C/C++ 代码、自定义 CUDA 内核或指向 C/C++/CUDA 扩展的 Python 绑定）时，才应使用算子注册 API 来创建自定义算子。

为什么要创建自定义算子？#

通过获取张量的数据指针并将其传递给 pybind 绑定后的内核，是可以实现使用 C/C++/CUDA 内核的。但是，这种方法无法与自动微分、torch.compile、vmap 等 PyTorch 子系统组合。为了让操作能与 PyTorch 子系统组合，它必须通过算子注册 API 进行注册。