自定义 C++ 和 CUDA 运算符

你将学到什么

如何将用 C++/CUDA 编写的自定义算子集成到 PyTorch 中
如何使用 torch.library.opcheck 测试自定义算子

先决条件

PyTorch 2.4 或更高版本
具备 C++ 和 CUDA 编程的基础知识

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

PyTorch 提供了大量用于操作张量的运算符库（例如 torch.add、torch.sum 等）。然而，您可能希望为 PyTorch 引入一个新的自定义运算符。本教程演示了使用 C++/CUDA 编写自定义运算符的最佳实践。

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

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

您可以在此找到本教程的完整工作示例 here。

配置构建系统

如果您正在开发自定义的 C++/CUDA 代码，这些代码必须被编译。请注意，如果您正在与一个已经绑定了预编译 C++/CUDA 代码的 Python 库进行交互，您可能会考虑编写一个自定义的 Python 操作符（参见自定义 Python 操作符）。

使用 torch.utils.cpp_extension 来编译自定义的 C++/CUDA 代码，以便与 PyTorch C++ 扩展一起使用。这些扩展可以通过 setuptools “提前”构建，或者通过 load_inline “即时”构建；我们将重点介绍“提前”构建的方式。

使用 cpp_extension 非常简单，只需编写如下 setup.py 文件：

fromsetuptoolsimport setup, Extension
fromtorch.utilsimport 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 包）。与 CPython 版本无关的特性有助于最大限度地减少您的自定义库需要支持和发布的 wheel 包数量。我们希望支持的最低版本是 3.9，因为它是目前支持的最旧版本，因此我们在整个设置代码中使用相应的十六进制编码和标识符。我们建议在您希望支持的最低 CPython 版本的环境中构建扩展，以最小化未知行为。因此，这里我们在 CPython 3.9 环境中构建扩展。构建完成后，这个单一的 wheel 包可以在任何 CPython 3.9+ 的环境中运行。为了实现这一点，有三行关键代码需要注意。

首先是在 extra_compile_args 中指定 Py_LIMITED_API，以支持最低版本的 CPython：

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

定义 Py_LIMITED_API 标志有助于验证扩展实际上仅使用了 CPython 稳定有限 API，这是构建 CPython 无关轮子的一个要求。如果未满足此要求，可能会构建一个看起来与 CPython 无关但在另一个 CPython 环境中会崩溃，或者更糟糕的是，会静默地出现错误的轮子。请务必避免使用不稳定的 CPython API，例如来自 libtorch_python 的 API（特别是 pytorch/python 绑定），并且仅使用来自 libtorch 的 API（ATen 对象、操作符和调度器）。我们强烈建议定义 Py_LIMITED_API 标志，以帮助确认扩展是合规的，并且作为 CPython 无关轮子是安全的。请注意，定义此标志并不能完全保证构建的轮子与 CPython 无关，但它比完全不加限制要好。在 Python 文档中提到了一些注意事项，您应该自己测试并验证轮子是否真正与相关 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
)

必须将 py_limited_api=True 作为参数传递给 CppExtension/CUDAExtension，并作为 "bdist_wheel" 命令的选项，同时指定支持的最低 CPython 版本（在本例中为 3.9）。因此，本教程中的 setup 将构建一个命名正确的 wheel 包，该包可以安装在多个 CPython 版本 >=3.9 上。

如果您的扩展使用了稳定有限集之外的 CPython API，那么您将无法构建一个与 CPython 版本无关的 wheel 包！您应该为每个 CPython 版本分别构建一个 wheel 包，如下所示：

fromsetuptoolsimport setup, Extension
fromtorch.utilsimport cpp_extension

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

定义自定义操作并添加后端实现

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

at::Tensormymuladd_cpu(at::Tensora,constat::Tensor&b,doublec){
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::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
at::Tensorresult=torch::empty(a_contig.sizes(),a_contig.options());
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();
for(int64_ti=0;i<result.numel();i++){
result_ptr[i]=a_ptr[i]*b_ptr[i]+c;
}
returnresult;
}

为了从 PyTorch 的 Python 前端使用它，我们需要使用 TORCH_LIBRARY API 将其注册为 PyTorch 操作符。这将自动将该操作符绑定到 Python。

操作符注册是一个两步过程：

定义运算符 - 此步骤确保 PyTorch 能够识别新的运算符。
注册后端实现 - 在此步骤中，将各种后端（如 CPU 和 CUDA）的实现与运算符关联起来。

定义操作符

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

为操作符选择一个命名空间。我们建议该命名空间为您的顶级项目名称；在本教程中，我们将使用“extension_cpp”。
提供一个模式字符串，用于指定操作符的输入/输出类型，以及输入张量是否会被修改。除了张量和浮点数之外，我们还支持更多类型；详情请参阅自定义操作符手册。
- 如果您正在编写一个可以修改其输入Tensors的运算符，请参阅此处（创建可变运算符）了解如何指定这一点。

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__voidmuladd_kernel(intnumel,constfloat*a,constfloat*b,floatc,float*result){
intidx=blockIdx.x*blockDim.x+threadIdx.x;
if(idx<numel)result[idx]=a[idx]*b[idx]+c;
}

at::Tensormymuladd_cuda(constat::Tensor&a,constat::Tensor&b,doublec){
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::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
at::Tensorresult=torch::empty(a_contig.sizes(),a_contig.options());
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();

intnumel=a_contig.numel();
muladd_kernel<<<(numel+255)/256,256>>>(numel,a_ptr,b_ptr,c,result_ptr);
returnresult;
}

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

为操作符添加 `torch.compile` 支持

要为操作符添加 torch.compile 支持，我们必须添加一个 FakeTensor 内核（也称为“元内核”或“抽象实现”）。FakeTensors 是具有元数据（如形状、数据类型、设备）但没有实际数据的张量：操作符的 FakeTensor 内核指定了如何根据输入张量的元数据计算输出张量的元数据。FakeTensor 内核应返回您选择的虚拟张量，并确保这些张量具有正确的元数据（形状/步幅/dtype/设备）。

我们建议通过 Python 的 torch.library.register_fake API 来实现这一点，尽管也可以通过 C++ 来实现（更多详细信息请参阅 The Custom Operators Manual）。

# 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。这可以通过以下三种方式实现：

加载包含自定义操作符定义的 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

在本教程中，由于我们重视能够跨多个 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)
{
staticstructPyModuleDefmodule_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 */
};
returnPyModule_Create(&module_def);
}
}

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

如果您希望在 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 文件的位置，以便加载它们，而这并不总是那么简单：

importtorch
frompathlibimport 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 或 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::Tensormymul_cpu(constat::Tensor&a,constat::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::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
at::Tensorresult=torch::empty(a_contig.sizes(),a_contig.options());
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();
for(int64_ti=0;i<result.numel();i++){
result_ptr[i]=a_ptr[i]*b_ptr[i];
}
returnresult;
}

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。

defsample_inputs(device, *, requires_grad=False):
    defmake_tensor(*size):
        return torch.randn(size, device=device, requires_grad=requires_grad)

    defmake_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],
    ]

defreference_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；否则，可能会出现未定义的行为。如果有多个变异的 Tensor，请为每个可变的 Tensor 使用不同的名称（例如，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.
voidmyadd_out_cpu(constat::Tensor&a,constat::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::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=out.data_ptr<float>();
for(int64_ti=0;i<out.numel();i++){
result_ptr[i]=a_ptr[i]+b_ptr[i];
}
}

在定义操作符时，我们必须在模式中指定它会改变输出张量：

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 的更多信息，请参阅自定义操作手册。