PyTorch自定义操作符¶

创建日期: 2024年6月18日 | 最后更新日期: 2024年7月22日 | 最后验证日期: 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 而不是 C++ 来编写自定义操作符，则可以考虑这样做：

您有一个Python函数希望PyTorch将其视为不透明的可调用对象，特别是在处理torch.compile和torch.export时。
您有一些Python绑定到C++/CUDA内核，并希望这些绑定能够与PyTorch子系统（如torch.compile或torch.autograd）组合使用

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

请参阅自定义 C++ 和 CUDA 操作符。

如果您希望从 C++ 而不是 Python 来编写自定义运算符，则可以考虑这样做：

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

自定义操作手册¶

对于本教程和页面未涵盖的信息，请参阅自定义操作手册（我们正在努力将这些信息移到我们的文档网站上）。我们建议您首先阅读上面的某个教程，然后将自定义操作手册作为参考；它并不是为了从头到尾逐字阅读。

我在什么情况下应该创建自定义操作符？¶

如果你的操作可以表示为 PyTorch 内置操作的组合，请将其编写为一个 Python 函数并调用它，而不是创建一个自定义操作。如果你需要调用的是 PyTorch 不理解的一些库（例如自定义的 C/C++ 代码、自定义的 CUDA 内核，或者连接到 C/C++/CUDA 扩展的 Python 绑定），则可以使用操作注册 API 创建一个自定义操作。

为什么应该创建自定义操作符？¶

可以使用 C/C++/CUDA 内核通过获取张量的数据指针并将其传递给 pybind 包装的内核来实现。然而，这种方法不支持与 PyTorch 子系统（如 autograd、torch.compile、vmap 等）的组合。为了使操作能够与 PyTorch 子系统组合，必须通过操作注册 API 进行注册。