PyTorch 自定义算子#

创建时间: 2024年6月18日 | 最后更新: 2025年7月31日 | 最后验证: 2024年11月5日

PyTorch 提供了大量可在 Tensor 上操作的算子（例如 torch.add, torch.sum 等）。然而，您可能希望将新的自定义操作引入 PyTorch，并使其能够与 torch.compile、autograd 和 torch.vmap 等子系统协同工作。要做到这一点，您必须通过 Python 的 torch.library 文档或 C++ 的 TORCH_LIBRARY API 将自定义操作注册到 PyTorch。

从 Python 编写自定义算子#

请参阅自定义 Python 算子。

如果您希望将一个 Python 函数作为不透明的可调用对象由 PyTorch 处理，特别是在 torch.compile 和 torch.export 方面，那么您可能希望从 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 是 Intel GPU 的后端编程语言。集成自定义 Sycl 代码请参考自定义 SYCL 算子。

如果您希望从 C++ 编写自定义算子（而不是 Python），那么您可能希望这样做，如果

您有自定义的 C++ 和/或 CUDA 代码。
您计划使用此代码与 AOTInductor 进行无 Python 推理。

自定义算子手册#

有关教程和本页面未涵盖的信息，请参阅自定义算子手册（我们正在努力将信息迁移到我们的文档网站）。我们建议您首先阅读以上任一教程，然后使用自定义算子手册作为参考；它不适合从头到尾阅读。

何时应该创建自定义算子？#

如果您的操作可以表示为内置 PyTorch 操作的组合，那么请将其写成一个 Python 函数并调用它，而不是创建自定义算子。如果您正在调用 PyTorch 无法识别的某个库（例如，自定义 C/C++ 代码、自定义 CUDA 内核或 C/C++/CUDA 扩展的 Python 绑定），请使用算子注册 API 来创建自定义算子。

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

可以通过获取 Tensor 的数据指针并将其传递给 pybind 内核来使用 C/C++/CUDA 内核。但是，这种方法无法与 autograd、torch.compile、vmap 等 PyTorch 子系统组合。为了使一个操作能够与 PyTorch 子系统组合，它必须通过算子注册 API 进行注册。