advanced/cpp_custom_ops

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

自定义 C++ 和 CUDA 算子

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

作者： Richard Zou

您将学到什么

• 如何将用 C++/CUDA 编写的自定义算子集成到 PyTorch 中

• 如何使用 torch.library.opcheck 测试自定义算子

先决条件

• PyTorch 2.4 或更高版本（若使用稳定 ABI，则为 PyTorch 2.10 或更高版本）

• 对 C++ 和 CUDA 编程有基本了解

注意

本教程无需任何额外修改即可在 AMD ROCm 上运行。

PyTorch 提供了大量的张量算子库（例如 torch.add, torch.sum 等）。但是，有时您可能希望将新的自定义算子引入 PyTorch。本教程演示了编写 C++/CUDA 自定义算子的推荐路径。

在本教程中，我们将演示如何编写一个融合乘加（fused multiply-add）的 C++ 和 CUDA 算子，使其能够与 PyTorch 子系统结合使用。该操作的语义如下

def mymuladd(a: Tensor, b: Tensor, c: float):
    return a * b + c

您可以在 extension-cpp 仓库中找到本教程的完整工作示例，其中包含两种并行实现

• extension_cpp/：使用标准 ATen/LibTorch API。

• extension_cpp_stable/：使用 LibTorch 稳定 ABI 支持的 API（推荐用于 PyTorch 2.10+）。

你应该使用哪个 API？

• ABI 稳定版 LibTorch API（推荐）：如果您使用的是 PyTorch 2.10+，我们建议使用 ABI 稳定版 API。它允许您构建一个能在多个 PyTorch 版本（2.10, 2.11, 2.12 等）下运行的 wheel 包，从而减轻支持多个 PyTorch 发行版的维护负担。有关更多详细信息，请参阅下方的 LibTorch 稳定 ABI（与 PyTorch 版本无关） 章节。

• 非 ABI 稳定版 LibTorch API：如果您需要稳定 ABI 中尚不可用的 API，或者您的目标是 2.10 之前的 PyTorch 版本，请使用此 API。请注意，您需要为您想要支持的每个 PyTorch 版本分别构建 wheel 包。

下方的代码片段通过标签页展示了两种实现方式，默认显示的是 ABI 稳定版 API。

设置构建系统

如果您正在开发自定义 C++/CUDA 代码，则必须对其进行编译。请注意，如果您正在对接一个已经绑定了预编译 C++/CUDA 代码的 Python 库，则可以考虑编写一个自定义 Python 算子（自定义 Python 算子）。

使用 torch.utils.cpp_extension 为 PyTorch 编译自定义 C++/CUDA 代码。C++ 扩展既可以通过 setuptools“预先”构建，也可以通过 load_inline“实时”构建；我们将重点介绍“预先”构建的方式。

使用 cpp_extension 就像编写一个 setup.py 一样简单

ABI 稳定版 LibTorch API

from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
              "extension_cpp",
              ["muladd.cpp"],
              extra_compile_args={
                  "cxx": [
                      # define Py_LIMITED_API with min version 3.9 to expose only the stable
                      # limited API subset from Python.h
                      "-DPy_LIMITED_API=0x03090000",
                      # define TORCH_TARGET_VERSION with min version 2.10 to expose only the
                      # stable API subset from torch
                      "-DTORCH_TARGET_VERSION=0x020a000000000000",
                  ]
              },
              py_limited_api=True)],  # Build 1 wheel across multiple Python versions
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      options={"bdist_wheel": {"py_limited_api": "cp39"}}  # 3.9 is minimum supported Python version
)

非 ABI 稳定版 LibTorch API

from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
              "extension_cpp",
              ["muladd.cpp"],
              extra_compile_args={
                  "cxx": [
                      "-DPy_LIMITED_API=0x03090000",
                  ]
              },
              py_limited_api=True)],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      options={"bdist_wheel": {"py_limited_api": "cp39"}}
)

如果您需要编译 CUDA 代码（例如 .cu 文件），请改用 torch.utils.cpp_extension.CUDAExtension。请参阅 extension-cpp 获取有关如何设置的示例。

与 CPython 版本无关性

上述示例代表了我们所说的“与 CPython 版本无关的 wheel 包”，这意味着我们构建了一个可以在多个 CPython 版本上运行的单一 wheel 包（类似于纯 Python 包）。为了尽量减少自定义库需要支持和发布的 wheel 包数量，这种无关性非常理想。我们希望支持的最低版本是 3.9，因为这是目前支持的最旧版本，因此我们在整个 setup 代码中使用了对应的十六进制代码和说明符。我们建议在您想要支持的最低 CPython 版本的环境中构建扩展，以最大限度地减少不可预知的行为。因此，在这里，我们在 CPython 3.9 环境中构建扩展。构建完成后，这个单一的 wheel 包将可以在任何 3.9+ 的 CPython 环境中运行。要实现这一点，有三个关键点需要注意。

第一点是在 extra_compile_args 中将 Py_LIMITED_API 指定为您想要支持的最低 CPython 版本

extra_compile_args={"cxx": ["-DPy_LIMITED_API=0x03090000"]},

定义 Py_LIMITED_API 标志有助于验证扩展是否确实仅使用了 CPython 稳定受限 API (CPython Stable Limited API)，这是构建与 CPython 无关的 wheel 包的要求。如果不满足此要求，则可能会构建出一个看起来与 CPython 无关但实际上会在另一个 CPython 环境中崩溃，甚至更糟糕的是，产生静默错误的 wheel 包。请务必避免使用不稳定的 CPython API（例如来自 libtorch_python 的 API，特别是 pytorch/python 绑定），并仅使用来自 libtorch 的 API（ATen 对象、算子和调度程序）。我们强烈建议定义 Py_LIMITED_API 标志，以帮助确定该扩展作为与 CPython 无关的 wheel 包是合规且安全的。请注意，定义此标志并不能完全保证构建的 wheel 包在所有 CPython 版本下都一定兼容，但它比完全没有规范要好得多。Python 文档中提到了一些注意事项，您应该亲自测试并验证该 wheel 包是否确实适用于相关的 CPython 版本。

第二行和第三行指定了 py_limited_api，告知 setuptools 您打算构建一个与 CPython 无关的 wheel 包，并将相应地影响 wheel 包的命名

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
            ...,
            py_limited_api=True)],  # Build 1 wheel across multiple Python versions
      ...,
      options={"bdist_wheel": {"py_limited_api": "cp39"}}  # 3.9 is minimum supported Python version
)

必须在 CppExtension/CUDAExtension 中指定 py_limited_api=True，并将其作为 "bdist_wheel" 命令的一个选项（带上最低支持的 CPython 版本，本例中为 3.9）。因此，我们教程中的 setup 将构建一个命名规范的 wheel 包，它可以安装在多个 >=3.9 的 CPython 版本中。

如果您的扩展使用了稳定受限集之外的 CPython API，则您无法构建与 CPython 无关的 wheel 包！在这种情况下，您应该为每个 CPython 版本构建一个单独的 wheel 包，如下所示

from setuptools import setup, Extension
from torch.utils import cpp_extension

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
            "extension_cpp",
            ["muladd.cpp"])],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
)

LibTorch 稳定 ABI（与 PyTorch 版本无关）

除了与 CPython 版本无关之外，还有第二个 wheel 包兼容性维度：与 LibTorch 版本无关。虽然与 CPython 无关允许构建一个在多个 Python 版本（3.9, 3.10, 3.11 等）下工作的单一 wheel 包，但与 LibTorch 无关则允许构建一个在多个 PyTorch 版本（2.10, 2.11, 2.12 等）下工作的单一 wheel 包。这两个概念是正交的，可以结合使用。

为了实现与 LibTorch 无关，您必须使用 ABI 稳定版 LibTorch API，它提供了一个用于与 PyTorch 张量和算子交互的稳定 API。例如，您不能再使用 at::Tensor，而必须使用 torch::stable::Tensor。有关稳定 ABI 的全面文档（包括迁移指南、支持的类型和基于栈的 API 约定），请参阅 LibTorch 稳定 ABI 文档。

稳定 ABI 的 setup.py 中包含了 TORCH_TARGET_VERSION=0x020a000000000000，这表示该扩展的目标是 LibTorch 稳定 ABI，最低支持的 PyTorch 版本为 2.10。版本格式为：[主版本号 1 字节][次版本号 1 字节][补丁版本 1 字节][ABI 标签 5 字节]，因此 2.10.0 = 0x020a000000000000。

如果稳定 API/ABI 不包含您需要的功能，您可以使用非 ABI 稳定版 LibTorch API，但您将需要为您想要支持的每个 PyTorch 版本构建单独的 wheel 包。

定义自定义算子并添加后端实现

首先，让我们编写一个计算 mymuladd 的 C++ 函数

ABI 稳定版 LibTorch API

#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/ops.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include <torch/headeronly/macros/Macros.h>

torch::stable::Tensor mymuladd_cpu(
    const torch::stable::Tensor& a,
    const torch::stable::Tensor& b,
    double c) {
  STD_TORCH_CHECK(a.sizes().equals(b.sizes()));
  STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(a.device().type() == torch::headeronly::DeviceType::CPU);
  STD_TORCH_CHECK(b.device().type() == torch::headeronly::DeviceType::CPU);

  torch::stable::Tensor a_contig = torch::stable::contiguous(a);
  torch::stable::Tensor b_contig = torch::stable::contiguous(b);
  torch::stable::Tensor result = torch::stable::empty_like(a_contig);

  const float* a_ptr = a_contig.const_data_ptr<float>();
  const float* b_ptr = b_contig.const_data_ptr<float>();
  float* result_ptr = result.mutable_data_ptr<float>();

  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = a_ptr[i] * b_ptr[i] + c;
  }
  return result;
}

非 ABI 稳定版 LibTorch API

#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>

at::Tensor mymuladd_cpu(const at::Tensor& a, const at::Tensor& b, double c) {
  TORCH_CHECK(a.sizes() == b.sizes());
  TORCH_CHECK(a.dtype() == at::kFloat);
  TORCH_CHECK(b.dtype() == at::kFloat);
  TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CPU);
  at::Tensor a_contig = a.contiguous();
  at::Tensor b_contig = b.contiguous();
  at::Tensor result = torch::empty(a_contig.sizes(), a_contig.options());
  const float* a_ptr = a_contig.data_ptr<float>();
  const float* b_ptr = b_contig.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();
  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = a_ptr[i] * b_ptr[i] + c;
  }
  return result;
}

为了从 PyTorch 的 Python 前端调用此函数，我们需要使用 TORCH_LIBRARY（或 STABLE_TORCH_LIBRARY）宏将其注册为 PyTorch 算子。这将自动把该算子绑定到 Python 中。

算子注册是一个两步过程

• 定义算子 - 此步骤确保 PyTorch 识别该新算子。

• 注册后端实现 - 在此步骤中，将针对 CPU 和 CUDA 等不同后端的实现与该算子关联起来。

定义算子

要定义算子，请遵循以下步骤

1. 为算子选择一个命名空间。我们建议命名空间与您的顶层项目名称一致；我们在本教程中使用“extension_cpp”。

2. 提供一个模式字符串（schema string），指定算子的输入/输出类型以及输入张量是否会被修改（mutated）。除了 Tensor 和 float 外，我们还支持更多类型；请参阅 自定义算子手册 以获取更多详细信息。

   • 如果您编写的算子会修改其输入张量，请参阅此处（创建可变算子）了解如何指定这一点。

ABI 稳定版 LibTorch API

STABLE_TORCH_LIBRARY(extension_cpp, m) {
  // Note that "float" in the schema corresponds to the C++ double type
  // and the Python float type.
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
}

非 ABI 稳定版 LibTorch API

TORCH_LIBRARY(extension_cpp, m) {
  // Note that "float" in the schema corresponds to the C++ double type
  // and the Python float type.
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
}

这使得该算子可以通过 torch.ops.extension_cpp.mymuladd 在 Python 中使用。

注册算子的后端实现

使用 TORCH_LIBRARY_IMPL（或 STABLE_TORCH_LIBRARY_IMPL）为算子注册后端实现。

ABI 稳定版 LibTorch API

请注意，我们使用 TORCH_BOX() 包装了函数指针——这是稳定 ABI 函数正确处理参数装箱/拆箱（boxing/unboxing）所必需的。

STABLE_TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", TORCH_BOX(&mymuladd_cpu));
}

非 ABI 稳定版 LibTorch API

TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", &mymuladd_cpu);
}

如果您也有 mymuladd 的 CUDA 实现，您可以在另一个 TORCH_LIBRARY_IMPL（或 STABLE_TORCH_LIBRARY_IMPL）块中注册它

ABI 稳定版 LibTorch API

#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/ops.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/csrc/stable/c/shim.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void muladd_kernel(int numel, const float* a, const float* b, float c, float* result) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < numel) result[idx] = a[idx] * b[idx] + c;
}

torch::stable::Tensor mymuladd_cuda(
    const torch::stable::Tensor& a,
    const torch::stable::Tensor& b,
    double c) {
  STD_TORCH_CHECK(a.sizes().equals(b.sizes()));
  STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(a.device().type() == torch::headeronly::DeviceType::CUDA);
  STD_TORCH_CHECK(b.device().type() == torch::headeronly::DeviceType::CUDA);

  torch::stable::Tensor a_contig = torch::stable::contiguous(a);
  torch::stable::Tensor b_contig = torch::stable::contiguous(b);
  torch::stable::Tensor result = torch::stable::empty_like(a_contig);

  const float* a_ptr = a_contig.const_data_ptr<float>();
  const float* b_ptr = b_contig.const_data_ptr<float>();
  float* result_ptr = result.mutable_data_ptr<float>();

  int numel = a_contig.numel();

  // For now, we rely on the raw shim API to get the current CUDA stream.
  // This will be improved in a future release.
  // When using a raw shim API, we need to use TORCH_ERROR_CODE_CHECK to
  // check the error code and throw an appropriate runtime_error otherwise.
  void* stream_ptr = nullptr;
  TORCH_ERROR_CODE_CHECK(
      aoti_torch_get_current_cuda_stream(a.get_device_index(), &stream_ptr));
  cudaStream_t stream = static_cast<cudaStream_t>(stream_ptr);

  muladd_kernel<<<(numel+255)/256, 256, 0, stream>>>(numel, a_ptr, b_ptr, c, result_ptr);
  return result;
}

STABLE_TORCH_LIBRARY_IMPL(extension_cpp, CUDA, m) {
  m.impl("mymuladd", TORCH_BOX(&mymuladd_cuda));
}

非 ABI 稳定版 LibTorch API

#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__global__ void muladd_kernel(int numel, const float* a, const float* b, float c, float* result) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < numel) result[idx] = a[idx] * b[idx] + c;
}

at::Tensor mymuladd_cuda(const at::Tensor& a, const at::Tensor& b, double c) {
  TORCH_CHECK(a.sizes() == b.sizes());
  TORCH_CHECK(a.dtype() == at::kFloat);
  TORCH_CHECK(b.dtype() == at::kFloat);
  TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CUDA);
  TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CUDA);
  at::Tensor a_contig = a.contiguous();
  at::Tensor b_contig = b.contiguous();
  at::Tensor result = at::empty(a_contig.sizes(), a_contig.options());
  const float* a_ptr = a_contig.data_ptr<float>();
  const float* b_ptr = b_contig.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();

  int numel = a_contig.numel();
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  muladd_kernel<<<(numel+255)/256, 256, 0, stream>>>(numel, a_ptr, b_ptr, c, result_ptr);
  return result;
}

TORCH_LIBRARY_IMPL(extension_cpp, CUDA, m) {
  m.impl("mymuladd", &mymuladd_cuda);
}

为算子添加 torch.compile 支持

要为算子添加 torch.compile 支持，我们必须添加一个 FakeTensor 内核（也称为“元内核”或“抽象实现”）。FakeTensor 是具有元数据（如形状、数据类型、设备）但没有实际数据的张量：算子的 FakeTensor 内核指定了如何在给定输入张量元数据的情况下计算输出张量的元数据。FakeTensor 内核应返回您选择的、具有正确张量元数据（形状/步长/dtype/设备）的占位张量。

我们建议通过 Python 使用 torch.library.register_fake API 来完成此操作，尽管也可以在 C++ 中完成（有关更多详细信息，请参阅 自定义算子手册）。

# Important: the C++ custom operator definitions should be loaded first
# before calling ``torch.library`` APIs that add registrations for the
# C++ custom operator(s). The following import loads our
# C++ custom operator definitions.
# Note that if you are striving for Python agnosticism, you should use
# the ``load_library(...)`` API call instead. See the next section for
# more details.
from . import _C

@torch.library.register_fake("extension_cpp::mymuladd")
def _(a, b, c):
    torch._check(a.shape == b.shape)
    torch._check(a.dtype == torch.float)
    torch._check(b.dtype == torch.float)
    torch._check(a.device == b.device)
    return torch.empty_like(a)

设置 Python/C++ 混合注册

在本教程中，我们在 C++ 中定义了一个自定义算子，在 C++ 中添加了 CPU/CUDA 实现，并在 Python 中添加了 FakeTensor 内核和反向传播公式。这些注册被加载（或导入）的顺序非常重要（顺序错误会导致错误）。

为了将自定义算子与 Python/C++ 混合注册一起使用，我们必须首先加载保存自定义算子定义的 C++ 库，然后调用 torch.library 注册 API。这可以通过以下三种方式实现

1. 加载包含自定义算子定义的 C++ 库的第一种方法是为 _C 定义一个虚拟的 Python 模块。然后，在 Python 中，当您使用 import _C 导入该模块时，与扩展对应的 .so 文件将被加载，并且 TORCH_LIBRARY 和 TORCH_LIBRARY_IMPL 静态初始化程序将会运行。您可以使用 PYBIND11_MODULE 创建一个虚拟 Python 模块（如下所示），但您会注意到它无法使用 Py_LIMITED_API 进行编译，因为 pybind11 并不保证仅使用稳定的受限 CPython API！使用下面的代码，遗憾的是您无法为您的扩展构建一个与 CPython 无关的 wheel 包！（预告：我想知道第二种方法是什么 ;) ）。

// in, say, not_agnostic/csrc/extension_BAD.cpp
#include <pybind11/pybind11.h>

PYBIND11_MODULE("_C", m) {}

# in, say, extension/__init__.py
from . import _C

2. 在本教程中，因为我们重视能够跨多个 CPython 版本构建单一 wheel 包的能力，我们将用稳定 API 调用替换不稳定的 PYBIND11 调用。下面的代码可以使用 -DPy_LIMITED_API=0x03090000 进行编译，并成功为我们的 _C 扩展创建了一个虚拟 Python 模块，以便它可以从 Python 中导入。有关更多详细信息，请参阅 extension_cpp/__init__.py 和 extension_cpp/csrc/muladd.cpp

#include <Python.h>

extern "C" {
  /* Creates a dummy empty _C module that can be imported from Python.
    The import from Python will load the .so consisting of this file
    in this extension, so that the TORCH_LIBRARY static initializers
    below are run. */
  PyObject* PyInit__C(void)
  {
      static struct PyModuleDef module_def = {
          PyModuleDef_HEAD_INIT,
          "_C",   /* name of module */
          NULL,   /* module documentation, may be NULL */
          -1,     /* size of per-interpreter state of the module,
                    or -1 if the module keeps state in global variables. */
          NULL,   /* methods */
      };
      return PyModule_Create(&module_def);
  }
}

# in, say, extension/__init__.py
from . import _C

3. 如果您想在 C++ 自定义算子中完全避免使用 Python.h，可以在 Python 中使用 torch.ops.load_library("/path/to/library.so") 来加载从扩展编译得到的 .so 文件。请注意，使用此方法时，不会为扩展创建 _C Python 模块，因此您不能在 Python 中调用 import _C。相反，torch.ops.load_library("/path/to/library.so") 将取代 import 语句来触发自定义算子的注册。接下来的挑战是了解 .so 文件的位置以便加载它们，这并不总是那么简单。

import torch
from pathlib import Path

so_files = list(Path(__file__).parent.glob("_C*.so"))
assert (
    len(so_files) == 1
), f"Expected one _C*.so file, found {len(so_files)}"
torch.ops.load_library(so_files[0])

from . import ops

为算子添加训练（自动微分）支持

使用 torch.library.register_autograd 为算子添加训练支持。优先使用此方法，而不是直接使用 Python 的 torch.autograd.Function（有关更多详细信息，请参阅 自定义算子手册）。

def _backward(ctx, grad):
    a, b = ctx.saved_tensors
    grad_a, grad_b = None, None
    if ctx.needs_input_grad[0]:
        grad_a = grad * b
    if ctx.needs_input_grad[1]:
        grad_b = grad * a
    return grad_a, grad_b, None

def _setup_context(ctx, inputs, output):
    a, b, c = inputs
    saved_a, saved_b = None, None
    if ctx.needs_input_grad[0]:
        saved_b = b
    if ctx.needs_input_grad[1]:
        saved_a = a
    ctx.save_for_backward(saved_a, saved_b)

# This code adds training support for the operator. You must provide us
# the backward formula for the operator and a `setup_context` function
# to save values to be used in the backward.
torch.library.register_autograd(
    "extension_cpp::mymuladd", _backward, setup_context=_setup_context)

请注意，反向传播必须是 PyTorch 可理解算子的组合。如果您希望在反向传播中使用另一个自定义 C++ 或 CUDA 内核，则必须将其封装为自定义算子。

如果我们有自己的自定义 mymul 内核，我们需要将其封装为自定义算子，然后从反向传播中调用它。

ABI 稳定版 LibTorch API

torch::stable::Tensor mymul_cpu(
    const torch::stable::Tensor& a,
    const torch::stable::Tensor& b) {
  STD_TORCH_CHECK(a.sizes().equals(b.sizes()));
  STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(a.device().type() == torch::headeronly::DeviceType::CPU);
  STD_TORCH_CHECK(b.device().type() == torch::headeronly::DeviceType::CPU);

  torch::stable::Tensor a_contig = torch::stable::contiguous(a);
  torch::stable::Tensor b_contig = torch::stable::contiguous(b);
  torch::stable::Tensor result = torch::stable::empty_like(a_contig);

  const float* a_ptr = a_contig.const_data_ptr<float>();
  const float* b_ptr = b_contig.const_data_ptr<float>();
  float* result_ptr = result.mutable_data_ptr<float>();

  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = a_ptr[i] * b_ptr[i];
  }
  return result;
}

STABLE_TORCH_LIBRARY(extension_cpp, m) {
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
  m.def("mymul(Tensor a, Tensor b) -> Tensor");
}

STABLE_TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", TORCH_BOX(&mymuladd_cpu));
  m.impl("mymul", TORCH_BOX(&mymul_cpu));
}

非 ABI 稳定版 LibTorch API

at::Tensor mymul_cpu(const at::Tensor& a, const at::Tensor& b) {
  TORCH_CHECK(a.sizes() == b.sizes());
  TORCH_CHECK(a.dtype() == at::kFloat);
  TORCH_CHECK(b.dtype() == at::kFloat);
  TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CPU);
  at::Tensor a_contig = a.contiguous();
  at::Tensor b_contig = b.contiguous();
  at::Tensor result = torch::empty(a_contig.sizes(), a_contig.options());
  const float* a_ptr = a_contig.data_ptr<float>();
  const float* b_ptr = b_contig.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();
  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = a_ptr[i] * b_ptr[i];
  }
  return result;
}

TORCH_LIBRARY(extension_cpp, m) {
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
  m.def("mymul(Tensor a, Tensor b) -> Tensor");
}

TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", &mymuladd_cpu);
  m.impl("mymul", &mymul_cpu);
}

def _backward(ctx, grad):
    a, b = ctx.saved_tensors
    grad_a, grad_b = None, None
    if ctx.needs_input_grad[0]:
        grad_a = torch.ops.extension_cpp.mymul.default(grad, b)
    if ctx.needs_input_grad[1]:
        grad_b = torch.ops.extension_cpp.mymul.default(grad, a)
    return grad_a, grad_b, None


def _setup_context(ctx, inputs, output):
    a, b, c = inputs
    saved_a, saved_b = None, None
    if ctx.needs_input_grad[0]:
        saved_b = b
    if ctx.needs_input_grad[1]:
        saved_a = a
    ctx.save_for_backward(saved_a, saved_b)


# This code adds training support for the operator. You must provide us
# the backward formula for the operator and a `setup_context` function
# to save values to be used in the backward.
torch.library.register_autograd(
    "extension_cpp::mymuladd", _backward, setup_context=_setup_context)

测试算子

使用 torch.library.opcheck 来测试自定义算子是否注册正确。请注意，此函数不测试梯度在数学上是否正确——请计划为其编写单独的测试，可以是手动测试，也可以使用 torch.autograd.gradcheck。

def sample_inputs(device, *, requires_grad=False):
    def make_tensor(*size):
        return torch.randn(size, device=device, requires_grad=requires_grad)

    def make_nondiff_tensor(*size):
        return torch.randn(size, device=device, requires_grad=False)

    return [
        [make_tensor(3), make_tensor(3), 1],
        [make_tensor(20), make_tensor(20), 3.14],
        [make_tensor(20), make_nondiff_tensor(20), -123],
        [make_nondiff_tensor(2, 3), make_tensor(2, 3), -0.3],
    ]

def reference_muladd(a, b, c):
    return a * b + c

samples = sample_inputs(device, requires_grad=True)
samples.extend(sample_inputs(device, requires_grad=False))
for args in samples:
    # Correctness test
    result = torch.ops.extension_cpp.mymuladd(*args)
    expected = reference_muladd(*args)
    torch.testing.assert_close(result, expected)

    # Use opcheck to check for incorrect usage of operator registration APIs
    torch.library.opcheck(torch.ops.extension_cpp.mymuladd.default, args)

创建可变算子

您可能希望编写一个会修改其输入的自定义算子。请在模式字符串中使用 Tensor(a!) 来指定每个可变张量；否则，将导致未定义的行为。如果有多个被修改的张量，请为每个可变张量使用不同的名称（例如 Tensor(a!), Tensor(b!), Tensor(c!)）。

让我们编写一个 myadd_out(a, b, out) 算子，它将 a+b 的内容写入 out。

ABI 稳定版 LibTorch API

void myadd_out_cpu(
    const torch::stable::Tensor& a,
    const torch::stable::Tensor& b,
    torch::stable::Tensor& out) {
  STD_TORCH_CHECK(a.sizes().equals(b.sizes()));
  STD_TORCH_CHECK(b.sizes().equals(out.sizes()));
  STD_TORCH_CHECK(a.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(b.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Float);
  STD_TORCH_CHECK(out.is_contiguous());
  STD_TORCH_CHECK(a.device().type() == torch::headeronly::DeviceType::CPU);
  STD_TORCH_CHECK(b.device().type() == torch::headeronly::DeviceType::CPU);
  STD_TORCH_CHECK(out.device().type() == torch::headeronly::DeviceType::CPU);

  torch::stable::Tensor a_contig = torch::stable::contiguous(a);
  torch::stable::Tensor b_contig = torch::stable::contiguous(b);

  const float* a_ptr = a_contig.const_data_ptr<float>();
  const float* b_ptr = b_contig.const_data_ptr<float>();
  float* result_ptr = out.mutable_data_ptr<float>();

  for (int64_t i = 0; i < out.numel(); i++) {
    result_ptr[i] = a_ptr[i] + b_ptr[i];
  }
}

在定义算子时，我们必须在模式字符串中指定它修改了 out 张量

STABLE_TORCH_LIBRARY(extension_cpp, m) {
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
  m.def("mymul(Tensor a, Tensor b) -> Tensor");
  m.def("myadd_out(Tensor a, Tensor b, Tensor(a!) out) -> ()");
}

STABLE_TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", TORCH_BOX(&mymuladd_cpu));
  m.impl("mymul", TORCH_BOX(&mymul_cpu));
  m.impl("myadd_out", TORCH_BOX(&myadd_out_cpu));
}

非 ABI 稳定版 LibTorch API

void myadd_out_cpu(const at::Tensor& a, const at::Tensor& b, at::Tensor& out) {
  TORCH_CHECK(a.sizes() == b.sizes());
  TORCH_CHECK(b.sizes() == out.sizes());
  TORCH_CHECK(a.dtype() == at::kFloat);
  TORCH_CHECK(b.dtype() == at::kFloat);
  TORCH_CHECK(out.dtype() == at::kFloat);
  TORCH_CHECK(out.is_contiguous());
  TORCH_INTERNAL_ASSERT(a.device().type() == at::DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(b.device().type() == at::DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(out.device().type() == at::DeviceType::CPU);
  at::Tensor a_contig = a.contiguous();
  at::Tensor b_contig = b.contiguous();
  const float* a_ptr = a_contig.data_ptr<float>();
  const float* b_ptr = b_contig.data_ptr<float>();
  float* result_ptr = out.data_ptr<float>();
  for (int64_t i = 0; i < out.numel(); i++) {
    result_ptr[i] = a_ptr[i] + b_ptr[i];
  }
}

在定义算子时，我们必须在模式字符串中指定它修改了 out 张量

TORCH_LIBRARY(extension_cpp, m) {
  m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
  m.def("mymul(Tensor a, Tensor b) -> Tensor");
  m.def("myadd_out(Tensor a, Tensor b, Tensor(a!) out) -> ()");
}

TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", &mymuladd_cpu);
  m.impl("mymul", &mymul_cpu);
  m.impl("myadd_out", &myadd_out_cpu);
}

注意

不要将任何被修改的张量作为算子的输出返回，因为这会导致与 torch.compile 等 PyTorch 子系统的不兼容。

结论

在本教程中，我们介绍了将自定义 C++ 和 CUDA 算子与 PyTorch 集成的推荐方法。TORCH_LIBRARY/STABLE_TORCH_LIBRARY 和 torch.library API 是相对底层的。有关如何使用该 API 的更多信息，请参阅 自定义算子手册。