评价此页
★ ★ ★ ★ ★

算子注册#

创建于：2025年8月27日 | 最后更新于：2025年9月2日

对于新的加速器，集成中最重要和最基本方面之一就是支持高性能算子。为了方便用户和加速器开发者进行算子适配，PyTorch 提供了多种方法来在 Python 和 C++ 中开发和注册算子。以下各节详细介绍了 PyTorch 在算子注册方面的一些基本功能。

注意

Dispatch Key 用于唯一标识 PyTorch 中的加速器，例如 CPU、CUDA、MPS 和 PrivateUse1。理论上，所有后续的新加速器都将共享 PrivateUse1，利用其内置的全面脚手架能力来完成新加速器的集成。如果您对 dispatcher 感兴趣，请参阅 Let’s talk about the PyTorch dispatcher。

算子集#

PyTorch 目前拥有超过 3500 个内置算子（包括相关的算子变体）。这无论从哪个角度来看都是一项巨大的工作量，而且在短时间内支持如此庞大的算子数量绝非易事。因此，作为开发新后端算子的第一步，我们的目标应该是专注于核心算子。对于其他算子，我们可以首先使用社区的 fallback 机制作为优先事项来支持该功能。之后，我们可以逐步完成其他算子，以提高新后端的性能。

所需的算子集列在下面，主要包括工厂函数所需的低级算子和 fallback 算子。

算子名称

Dispatch Key

描述

empty.memory_format

PrivateUse1

使用指定的形状和内存布局（步幅自动计算）创建未初始化的 Tensor。

empty_strided

PrivateUse1

创建具有指定形状和步幅的未初始化 Tensor（具有更大的自由度）。

as_strided

PrivateUse1

使用新的形状、步幅和偏移量创建输入 Tensor 的共享视图（无需分配新内存）。

view

PrivateUse1

创建具有新形状的输入 Tensor 的共享视图，但原始 Tensor 必须是内存连续的。

_reshape_alias

PrivateUse1

创建无安全检查的共享视图（reshape 的内部版本）。

resize_

PrivateUse1

就地修改 Tensor 的形状，并在容量不足时重新分配内存。

_copy_from

PrivateUse1

Tensor.copy_ 的底层核心函数，负责实际的跨设备数据复制。

_copy_from_and_resize

PrivateUse1

结合 resize_ 和 _copy_from，先调整大小再复制。

_local_scalar_dense

PrivateUse1

.item() 的底层实现，将 Tensor 中的值提取为 CPU 标量。

set_.source_Tensor

PrivateUse1

使用指定的 Tensor 设置当前 Tensor。

set_.source_Storage

PrivateUse1

使用指定的 Storage 设置当前 Tensor。

set_.source_Storage_storage_offset

PrivateUse1

使用指定的 Storage 和存储偏移量设置当前 Tensor。

fallback

PrivateUse1

回退到 CPU。

基础#

现在我们已经定义了算子支持的初始范围，我们可以开始开发算子适配。本节将根据实际场景，在 Python 和 C++ 中解释这些实现。

第一步#

上面提到的算子 具有一个共同点：它们是内置的 PyTorch 算子，具有定义的 命名空间 和 Schema，并且这些算子的内置加速器（CPU、CUDA 等）已经实现。我们接下来要做的是为新加速器实现这些算子。

 1at::Tensor empty_memory_format(
 2    c10::IntArrayRef size,
 3    std::optional<c10::ScalarType> dtype_opt,
 4    std::optional<c10::Layout> layout_opt,
 5    std::optional<c10::Device> device_opt,
 6    std::optional<bool> pin_memory_opt,
 7    std::optional<c10::MemoryFormat> memory_format_opt) {
 8  const auto device = c10::device_or_default(device_opt);
 9  const auto dtype = c10::dtype_or_default(dtype_opt);
10  TORCH_CHECK(device.is_privateuseone());
11  TORCH_CHECK(
12      c10::layout_or_default(layout_opt) == c10::Layout::Strided,
13      "Non strided layout not supported");
14  TORCH_CHECK(
15      !c10::pinned_memory_or_default(pin_memory_opt),
16      "Pin memory can only be on CPU");
17  const c10::DeviceGuard device_guard(device);
18  constexpr c10::DispatchKeySet pu1_dks(c10::DispatchKey::PrivateUse1);
19  auto allocator = at::GetAllocator(at::kPrivateUse1);
20  return at::detail::empty_generic(
21      size, allocator, pu1_dks, dtype, memory_format_opt);
22}

 1at::Tensor wrapper_empty_memory_format(
 2    c10::IntArrayRef size,
 3    std::optional<c10::ScalarType> dtype_opt,
 4    std::optional<c10::Layout> layout_opt,
 5    std::optional<c10::Device> device_opt,
 6    std::optional<bool> pin_memory_opt,
 7    std::optional<c10::MemoryFormat> memory_format_opt) {
 8  return at::native::openreg::empty_memory_format(
 9      size,
10      dtype_opt,
11      layout_opt,
12      device_opt,
13      pin_memory_opt,
14      memory_format_opt);
15}

以 empty.memory_format 算子为例，我们首先需要在 native_functions.yaml 中查询算子的 schema 信息，其中包含详细的签名信息。然后，我们可以根据新加速器的能力来实现该算子。

- func: empty.memory_format(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
dispatch:
    CPU: empty_cpu
    CUDA: empty_cuda
    ...

 1TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
 2  m.impl("empty.memory_format", wrapper_empty_memory_format);
 3  m.impl("empty_strided", wrapper_empty_strided);
 4  m.impl("as_strided", wrapper_as_strided);
 5  m.impl("resize_", wrapper_resize_);
 6  m.impl("_reshape_alias", wrapper__reshape_alias);
 7  m.impl("_copy_from", wrapper__copy_from);
 8  m.impl("_copy_from_and_resize", wrapper__copy_from_and_resize);
 9  m.impl("_local_scalar_dense", wrapper__local_scalar_densor);
10  m.impl("set_.source_Tensor", wrapper_set_source_Tensor_);
11  m.impl("set_.source_Storage", wrapper_set_source_Storage_);
12  m.impl(
13      "set_.source_Storage_storage_offset",
14      wrapper_set_source_Storage_storage_offsetset_);
15  m.impl("view", wrapper_view);
16}

完成 wrapper_empty_memory_format 后，我们可以通过 TORCH_LIBRARY_IMPL 为 PrivateUse1 注册 aten::empty.memory_format。

第二步#

按照 第一步，我们可以完成除 fallback 之外所有算子的开发和注册。接下来，为了支持与运算相关的算子（例如数学运算和卷积运算），我们需要实现 fallback 语义的注册。这是 PyTorch 框架提供的内置功能，可以将新加速器不支持的某些运算回退到 CPU 执行。对于正在开发的新后端，这是确保功能性的极其有效的方法，但会牺牲性能。

 1void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
 2  static const std::unordered_set<c10::OperatorName> cpu_fallback_blocklist = {
 3      c10::OperatorName("aten::abs", ""),
 4      c10::OperatorName("aten::abs", "out"),
 5  };
 6
 7  const auto& op_name = op.schema().operator_name();
 8  if (cpu_fallback_blocklist.count(op_name)) {
 9    TORCH_CHECK(
10        false,
11        "Operator '",
12        op_name,
13        "' is not implemented for device openreg.");
14  } else {
15    at::native::cpu_fallback(op, stack);
16  }
17}

1void wrapper_cpu_fallback(
2    const c10::OperatorHandle& op,
3    torch::jit::Stack* stack) {
4  at::native::openreg::cpu_fallback(op, stack);
5}

1TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
2  m.fallback(
3      torch::CppFunction::makeFromBoxedFunction<&wrapper_cpu_fallback>());
4}

wrapper_cpu_fallback 封装了 PyTorch 提供的 at::native::cpu_fallback 方法，并通过 TORCH_LIBRARY_IMPL 在 PyTorch 中注册到 PrivateUse1。后续新后端不支持的操作将自动回退到 CPU 执行，执行完成后结果将传回新后端。

高级#

选择性回退#

仅为某些算子启用回退机制，而其他算子则遵循 PyTorch 的默认行为（如果加速器没有相应的算子实现，则会报错），这是一种非常合理的场景。

1void wrapper_cpu_fallback(
2    const c10::OperatorHandle& op,
3    torch::jit::Stack* stack) {
4  at::native::openreg::cpu_fallback(op, stack);
5}

1TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
2  m.impl(
3      "sub.Tensor",
4      torch::CppFunction::makeFromBoxedFunction<&wrapper_cpu_fallback>());
5}

每个算子的回退与全局回退非常相似，唯一的区别在于注册方法：调用 m.impl 为特定算子注册实现，而 m.fallback 为所有算子注册默认实现。

 1void cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
 2  static const std::unordered_set<c10::OperatorName> cpu_fallback_blocklist = {
 3      c10::OperatorName("aten::abs", ""),
 4      c10::OperatorName("aten::abs", "out"),
 5  };
 6
 7  const auto& op_name = op.schema().operator_name();
 8  if (cpu_fallback_blocklist.count(op_name)) {
 9    TORCH_CHECK(
10        false,
11        "Operator '",
12        op_name,
13        "' is not implemented for device openreg.");
14  } else {
15    at::native::cpu_fallback(op, stack);
16  }
17}

当然，全局回退也可以与回退黑名单结合使用，这是一种常见的方法，尤其是在只有少数算子不支持回退的情况下。

PyTorch STUB#

PyTorch 还为内置算子提供了另一种方法：STUB。此方法本质上基于 第一步<step-one> 方法，但增加了二次调度功能（例如，基于 CPU 特性的调度）。

注意

STUB 方法目前仅支持有限的算子集。对于新的加速器设备，STUB 方法的优势在于它以少量的性能开销为代价，大大降低了开发成本。PyTorch 目前没有明确列出可以通过 STUB 注册的算子集。由于相关算子数量庞大，此处仅提供支持算子列表的查询方法。

pushd ${TORCH_ROOT}

find aten -type f -a -name "*.h" | xargs -I {} grep -wl "^DECLARE_DISPATCH" {}

popd

DECLARE_DISPATCH 是一个宏，用于显式声明 STUB。它目前分布在 aten 目录中。基于此宏，您可以找到所有可以使用 STUB 方法集成的算子。

...
aten/src/ATen/native/Activation.h
aten/src/ATen/native/FusedSGD.h
aten/src/ATen/native/nested/NestedTensorBinaryOps.h
aten/src/ATen/native/TensorCompare.h
aten/src/ATen/native/Sorting.h
...

using unary_fn = void(*)(TensorIteratorBase&);

DECLARE_DISPATCH(unary_fn, abs_stub)

上面的列表包含声明 STUB 算子的文件，您可以在其中清楚地看到 STUB 名称和相关的函数签名。接下来，我们将以 abs_stub 为例，简要介绍通过 STUB 支持算子的路径。

 1void abs_kernel(at::TensorIteratorBase& iter) {
 2  TORCH_CHECK(iter.ntensors() == 2, "Abs kernel expects 2 tensors");
 3  TORCH_CHECK(
 4      iter.common_dtype() == at::ScalarType::Float,
 5      "Abs kernel only supports float type");
 6
 7  auto& output_tensor = iter.tensor(0);
 8  auto& input_tensor = iter.tensor(1);
 9
10  TORCH_CHECK(
11      input_tensor.sizes() == output_tensor.sizes(),
12      "Input and output tensor sizes must match.");
13
14  auto abs_loop = [](float* out_ptr, const float* in_ptr, int64_t n) {
15    for (int64_t i = 0; i < n; ++i) {
16      out_ptr[i] = std::abs(in_ptr[i]);
17    }
18  };
19
20  MemoryGuard guard(input_tensor, output_tensor);
21
22  if (iter.is_contiguous()) {
23    abs_loop(
24        static_cast<float*>(iter.data_ptr(0)),
25        static_cast<float*>(iter.data_ptr(1)),
26        iter.numel());
27  } else {
28    TORCH_CHECK(
29        input_tensor.is_contiguous(), "Input tensor must be contiguous.")
30
31    auto output = at::empty(
32        input_tensor.sizes(),
33        input_tensor.options().memory_format(
34            input_tensor.suggest_memory_format()));
35
36    MemoryGuard guard(output);
37
38    abs_loop(
39        static_cast<float*>(output.data_ptr()),
40        static_cast<float*>(iter.data_ptr(1)),
41        iter.numel());
42
43    output_tensor.copy_(output);
44  }
45}

1REGISTER_PRIVATEUSE1_DISPATCH(abs_stub, &wrapper_abs_stub);
2REGISTER_PRIVATEUSE1_DISPATCH(
3    quantize_tensor_per_tensor_affine_stub,
4    &wrapper_quantize_tensor_per_tensor_affine_stub);
5REGISTER_PRIVATEUSE1_DISPATCH(
6    _fused_sdp_choice_stub,
7    &wrapper__fused_sdp_choice);

从签名可以看出，abs_stub 的输入是 TensorIteratorBase，这是 PyTorch 提供的一个强大的辅助类，包含所有输入和输出算子以及一些其他辅助方法。基于此，我们可以开发 abs_kernel 算子，然后调用 REGISTER_PRIVATEUSE1_DISPATCH 来指定 abs_stub 以完成注册。

自定义算子#

除了 PyTorch 的内置算子，自定义加速器算子在提高特定场景下的性能方面也非常常见。这些可以分为三种主要方法：

  • 仅前向传播

  • 前向和后向传播：分开注册

  • 前向和后向传播：使用 torch.autograd.Function 实现

注意

PyTorch 教程中有更多细节，如果您有兴趣，请参考 PyTorch Custom Operators。

仅前向传播#

这里，我们将简要介绍自定义算子的实现过程，重点关注仅前向传播的方法。实现可以概括为以下三点：

  1. 定义 Schema

    1TORCH_LIBRARY(openreg, m) {
2  m.def("custom_abs(Tensor input)-> Tensor");
3}

    • 命名空间名称：openreg

    • 函数名称：custom_abs

    • 输入参数

      • 类型：Tensor

      • 名称：input

    • 输出类型：Tensor

  2. 注册算子与自动求导回退

    1TORCH_LIBRARY_IMPL(openreg, PrivateUse1, m) {
2  m.impl("custom_abs", &wrapper_custom_abs);
3}

    1TORCH_LIBRARY_IMPL(_, AutogradPrivateUse1, m) {
2  m.fallback(torch::autograd::autogradNotImplementedFallback());
3}

    使用 TORCH_LIBRARY_IMPL 为 PrivateUse1 中的 custom_abs 算子注册 wrapper_custom_abs 实现。但是，由于 PyTorch 中始终启用 Autograd，即使只需要前向计算，PyTorch 也会默认查找并执行相应的后向实现（将在后向实现中 fallthrough）。因此，我们还需要为 custom_abs 算子的 AutogradPrivateUse1 注册相应的实现。幸运的是，PyTorch 还提供了一个通用的 Autograd Fallback 机制，名为 torch::autograd::autogradNotImplementedFallback，如果仅涉及前向计算，它相当于一个 fallthrough 操作，选择下一个 DispatchKey 进行计算；如果涉及后向计算，则会抛出错误。

  3. 注册元数据（可选，但图模式等需要）

    1lib = torch.library.Library("openreg", "IMPL", "Meta")  # noqa: TOR901
2
3
4@torch.library.impl(lib, "custom_abs")
5def custom_abs(self):
6    return torch.empty_like(self)
7
8

    PyTorch 支持在 C++ 和 Python 中注册 Meta。由于 Python 注册更简单，因此此处以 Python 为例。与 C++ 中的 TORCH_LIBRARY_IMPL 函数类似，Python 提供了更友好的 torch.library.impl 装饰器。

工具#

PyTorch 中的算子注册很复杂，注册方法多样，场景众多。因此，PyTorch 社区提供了一些工具来帮助开发者快速理解底层原理并协助故障排除。在此我们简要介绍几种常用工具：

命令#

PyTorch 提供了围绕其 Dispatch 功能的一系列以 torch._C._dispatch_ 开头的命令。您可以使用以下命令查询所有相关的接口。

python -c 'import torch; print("\n".join([x for x in dir(torch._C) if x.startswith("_dispatch_")]))'

...
_dispatch_dump
_dispatch_dump_table
_dispatch_has_kernel
_dispatch_has_kernel_for_any_dispatch_key
_dispatch_has_kernel_for_dispatch_key
_dispatch_isTensorSubclassLike
_dispatch_is_alias_key
_dispatch_is_included_in_alias
_dispatch_is_main_interpreter
_dispatch_kernel_for_dispatch_key_is_fallthrough
_dispatch_key_for_device
_dispatch_key_name
_dispatch_key_parse
_dispatch_key_set
...

以下是几个常用命令的解释：

  • torch._C._dispatch_key_set:

    显示当前 Tensor 的 DispatchKey，优先级从左到右递增。

    >>> import torch
>>> a = torch.randn(3,3,device="cuda")
>>> torch._C._dispatch_key_set(a)
'DispatchKeySet(CUDA, ADInplaceOrView, AutogradCUDA, AutocastCUDA)'

  • torch._C._dispatch_dump_table:

    查询给定算子在不同 Dispatch Keys 下的支持状态，方便定位对应的实现代码。

    >>> import torch
>>> print(torch._C._dispatch_dump_table("aten::add.Tensor"))
>>> ...
    CPU: registered at ./build/aten/src/ATen/RegisterCPU_0.cpp:1309 [kernel]
    CUDA: registered at ./build/aten/src/ATen/RegisterCUDA_0.cpp:2420 [kernel]
    HIP: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    MPS: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    IPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    XPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    HPU: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    VE: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    MTIA: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    MAIA: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    PrivateUse1: registered at ./build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional_0.cpp:1373 [default backend kernel]
    ...

    您可以轻松地查询其他平台上 aten::add.Tensor 算子的对应实现，从而可以从源代码级别跟踪整个算子调用过程。

环境变量#

PyTorch 还提供了一些与 dispatcher 相关的环境变量，有助于学习和快速定位问题。

  • TORCH_SHOW_DISPATCH_TRACE

    显示 PyTorch 执行过程中的详细内部 dispatch key 调度。

    export TORCH_SHOW_DISPATCH_TRACE=1

    >>> import torch
>>> a = torch.randn(3,3)
 [call] op=[aten::randn], key=[BackendSelect]
   [redispatch] op=[aten::randn], key=[CPU]
     [call] op=[aten::empty.memory_format], key=[BackendSelect]
       [redispatch] op=[aten::empty.memory_format], key=[CPU]
     [call] op=[aten::normal_], key=[CPU]

    您可以清楚地看到 PyTorch 中 Python 级算子调用的所有底层算子：包括算子名称、调用层级以及对应的 Dispatch Key。

在此页面上
在 GitHub 上编辑
显示源代码
PyTorch 库