目录
  • 指南 >
  • 自定义 C++和 CUDA 算子
快捷键
高级/C++自定义操作

自定义 C++和 CUDA 算子

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：邹理

你将学到什么

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

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

前提条件

  • PyTorch 2.4 或更高版本

  • C++和 CUDA 编程的基本理解

备注

本教程也适用于 AMD ROCm，无需任何额外修改。

PyTorch 提供了一组大型算子库，这些算子可以在张量上工作（例如 torch.add、torch.sum 等）。然而，您可能希望将一个新的自定义算子引入 PyTorch。本教程演示了在 C++/CUDA 中编写自定义算子的官方途径。

在我们的教程中，我们将演示如何编写一个融合乘加的 C++和 CUDA 算子，该算子可以与 PyTorch 子系统协同工作。该操作的语义如下：

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

您可以在此处找到本教程的端到端工作示例。

设置构建系统

如果您正在开发自定义的 C++/CUDA 代码，则必须对其进行编译。请注意，如果您正在与已经具有对预编译 C++/CUDA 代码的绑定的 Python 库进行交互，您可以考虑编写自定义 Python 运算符（自定义 Python 运算符）。

使用 torch.utils.cpp_extension 编译自定义 C++/CUDA 代码以与 PyTorch 一起使用，C++扩展可以“预编译”方式通过 setuptools 构建，或者通过 load_inline“即时”构建；我们将重点关注“预编译”方式。

使用 cpp_extension 就像编写以下 setup.py 一样简单：

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

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
            "extension_cpp",
            ["muladd.cpp"],
            # define Py_LIMITED_API with min version 3.9 to expose only the stable
            # limited API subset from Python.h
            extra_compile_args={"cxx": ["-DPy_LIMITED_API=0x03090000"]},
            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
)

如果您需要编译 CUDA 代码（例如， .cu 文件），则请使用 torch.utils.cpp_extension.CUDAExtension。请参阅 extension-cpp 了解如何设置示例。

上面的示例代表我们所说的 CPython 兼容的 wheel，这意味着我们正在构建一个可以在多个 CPython 版本上运行的单个 wheel（类似于纯 Python 包）。在最小化您自定义库需要支持和支持的 wheel 数量方面，CPython 兼容性是可取的。我们希望支持的最小版本是 3.9，因为它是目前支持的最老版本，所以我们使用相应的 hexcode 和指定符贯穿整个设置代码。我们建议在您希望支持的最低 CPython 版本的环境中构建扩展，以最小化未知行为，因此，在这里我们在 CPython 3.9 环境中构建扩展。构建完成后，这个单个 wheel 将可以在任何 CPython 3.9+ 环境中运行。为了实现这一点，有三个关键行需要注意。

第一条是 Py_LIMITED_API 在 extra_compile_args 中的指定，以支持您希望的最小 CPython 版本：

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

定义 Py_LIMITED_API 标志有助于验证扩展实际上只使用了 CPython 稳定有限 API，这是构建 CPython 无关 wheel 的要求。如果这个要求没有得到满足，就有可能构建出一个看似 CPython 无关但实际上会在其他 CPython 环境中崩溃，或者更糟糕的是，会静默错误的 wheel。请小心避免使用不稳定的 CPython API，例如来自 libtorch_python（特别是 pytorch/python 绑定）的 API，而只使用来自 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，可以在多个 CPython 版本上安装 >=3.9 。

如果您的扩展使用 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},
)

定义自定义操作符和添加后端实现

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

at::Tensor mymuladd_cpu(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 API 将其注册为 PyTorch 运算符。这将自动将运算符绑定到 Python。

运算符注册是一个两步过程：

  • 定义运算符 - 此步骤确保 PyTorch 了解新的运算符。

  • 注册后端实现 - 在此步骤中，将 CPU 和 CUDA 等各种后端的实现与运算符关联起来。

定义一个操作符 ¶

要定义一个操作符，请按照以下步骤操作：

  1. 选择一个操作符的命名空间。我们建议命名空间使用顶级项目的名称；在我们的教程中我们将使用“extension_cpp”。

  2. 提供一个模式字符串，该字符串指定操作符的输入/输出类型以及输入 Tensors 是否会被修改。除了 Tensor 和 float 类型外，我们还支持更多类型；请参阅《自定义操作符手册》以获取更多详细信息。

    • 如果您正在编写可以修改其输入张量的运算符，请参阅此处（创建可变运算符）了解如何指定。

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 注册操作符的后端实现。

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

如果您也有 myaddmul 的 CUDA 实现，您可以在一个单独的 TORCH_LIBRARY_IMPL 块中注册它：

__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 = 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>();

  int numel = a_contig.numel();
  muladd_kernel<<<(numel+255)/256, 256>>>(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 内核（也称为“元内核”或“抽象实现”）。FakeTensors 是具有元数据（如形状、数据类型、设备）但没有数据的张量：操作符的 FakeTensor 内核指定了如何根据输入张量的元数据计算输出张量的元数据。FakeTensor 内核应返回您选择的虚拟张量，并具有正确的张量元数据（形状/步长/ dtype /设备）。

我们建议通过 torch.library.register_fake API 从 Python 执行此操作，尽管也可以从 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 静态初始化器。可以通过以下方式创建一个虚拟 Python 模块 PYBIND11_MODULE ，但你将注意到，这不能与 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

  1. 在本教程中，因为我们重视能够在多个 CPython 版本之间构建单个轮子，所以我们将用稳定的 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

  1. 如果您想在您的 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") 可以做到这一点。那么，挑战就转移到了理解 .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

为运算符添加训练（autograd）支持

使用 torch.library.register_autograd 为运算符添加训练支持。优先使用此方法，而不是直接使用 Python 的 torch.autograd.Function 或 C++的 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 内核，则需要将其封装成自定义运算符，然后从反向传播中调用它：

// New! a mymul_cpu kernel
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_CHECK(a.device().type() == at::DeviceType::CPU);
  TORCH_CHECK(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");
  // New! defining the mymul operator
  m.def("mymul(Tensor a, Tensor b) -> Tensor");
}


TORCH_LIBRARY_IMPL(extension_cpp, CPU, m) {
  m.impl("mymuladd", &mymuladd_cpu);
  // New! registering the cpu kernel for the mymul operator
  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 。

// An example of an operator that mutates one of its inputs.
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");
  // New!
  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);
  // New!
  m.impl("myadd_out", &myadd_out_cpu);
}

备注

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

结论 ¶

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

评价这个教程

© 版权所有 2024，PyTorch。

使用 Sphinx 构建，主题由 Read the Docs 提供。
//暂时添加调查链接

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookies。通过点击或导航，您同意允许我们使用 cookies。作为本站点的当前维护者，适用 Facebook 的 Cookies 政策。了解更多信息，包括可用的控制选项：Cookies 政策。