PyTorch 自定义算子

PyTorch 提供了大量用于张量运算的算子库（例如 torch.add、torch.sum 等）。但是，你可能希望将新的自定义操作引入 PyTorch，并使其与 torch.compile、自动微分和 torch.vmap 等子系统一起使用。为此，必须通过 Python torch.library 文档 或 C++ TORCH_LIBRARY API 将自定义操作注册到 PyTorch。

从 Python 编写自定义算子

请参阅 Python 自定义算子.

你可能希望从 Python（而不是 C++）编写自定义算子，如果

• 你有一个 Python 函数，希望 PyTorch 将其视为不透明的可调用对象，尤其是在 torch.compile 和 torch.export 方面。

• 你有一些指向 C++/CUDA 内核的 Python 绑定，并希望这些绑定与 PyTorch 子系统（如 torch.compile 或 torch.autograd）进行组合。

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

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

你可能希望从 C++（而不是 Python）编写自定义算子，如果

• 你拥有自定义 C++ 和/或 CUDA 代码。

• 你计划将此代码与 AOTInductor 一起使用以进行无 Python 推理。

自定义算子手册

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

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

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

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

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