自定义C++和CUDA运算符¶

创建日期: 2024年6月18日 | 最后更新日期: 2024年12月30日 | 最后验证日期: 2024年11月5日

你将学到什么

如何将用C++/CUDA编写的自定义操作与PyTorch集成
如何测试自定义运算符使用 torch.library.opcheck

先决条件

PyTorch 2.4 或更高版本
基本掌握C++和CUDA编程

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 来编译与 PyTorch 一起使用的自定义 C++/CUDA 代码 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"])],
      cmdclass={'build_ext': cpp_extension.BuildExtension})

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

从PyTorch 2.6开始，您现在可以为多个CPython版本构建单个wheel（类似于纯Python包的操作）。特别是，如果您的自定义库遵循CPython稳定有限API或完全避免使用CPython，您可以通过setuptools的py_limited_api标志针对支持的最低CPython版本构建一个与Python无关的wheel，如下所示：

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

setup(name="extension_cpp",
      ext_modules=[
          cpp_extension.CppExtension(
            "extension_cpp",
            ["python_agnostic_code.cpp"],
            py_limited_api=True)],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      options={"bdist_wheel": {"py_limited_api": "cp39"}}
)

请注意，您必须在 py_limited_api=True 之内指定 setup，同时也要作为 "bdist_wheel" 命令的选项之一，使用最小支持的 Python 版本（在这种情况下为 3.9）。这 setup 将构建一个可以在多个 Python 版本之间安装的单个 wheel 文件。python>=3.9 请参见 torchao 以获取示例。

注意

您必须独立验证构建的wheel文件是否真正兼容所有Python版本。指定py_limited_api并不会提供任何保证，因此有可能构建一个看似兼容所有Python环境的wheel文件，但在另一个Python环境中会崩溃，或者更糟糕的是，会默默地给出错误的结果。请务必避免使用不稳定的CPython API，例如libtorch_python中的API（特别是pytorch/python绑定），并仅使用libtorch中的API（aten对象、操作符和调度器）。例如，为了从Python访问自定义操作，库应该通过调度器注册这些操作符（详见下方！）。

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

首先，让我们编写一个计算 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）的实现与操作关联起来。

定义一个操作符¶

要定义一个操作符，请遵循这些步骤：

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

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);
}

如果您也有 CUDA 实现的 myaddmul，您可以在单独的 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++中添加了CPU/CUDA实现，并在Python中添加了FakeTensor核函数和反向公式。这些注册的加载顺序（或导入顺序）很重要（导入顺序错误会导致错误）。

要使用带有混合Python/C++注册的自定义操作符，我们必须首先加载包含自定义操作符定义的C++库，然后调用torch.library注册API。这可以有以下两种方式：

如果你正在跟随这个教程，导入我们创建的Python C扩展模块将会加载自定义操作符的C++定义。
如果您自定义的C++运算符位于共享库对象中，您也可以使用torch.ops.load_library("/path/to/library.so")来加载它。这是实现Python独立性的官方途径，因为您不需要导入Python C扩展模块。请参见torchao __init__.py 以获取示例。

为操作符添加训练（自动求导）支持¶

使用 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) 操作符，which 写入 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的更多信息，请参阅自定义操作符手册。

自定义C++和CUDA运算符¶

设置构建系统¶

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

定义一个操作符¶

为运算符注册后端实现¶

添加 `torch.compile` 支持给一个操作符¶

设置混合Python/C++注册¶

为操作符添加训练（自动求导）支持¶

测试一个操作符¶

创建可变操作符¶

结论¶

文档

教程

资源

自定义C++和CUDA运算符¶

设置构建系统¶

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

定义一个操作符¶

为运算符注册后端实现¶

添加 torch.compile 支持给一个操作符¶

设置混合Python/C++注册¶

为操作符添加训练（自动求导）支持¶

测试一个操作符¶

创建可变操作符¶

结论¶

文档

教程

资源

添加 `torch.compile` 支持给一个操作符¶