在 C++ 中注册一个 Dispatched Operator¶

创建于: 2020 年 7 月 22 日 | 最后更新于: 2024 年 7 月 22 日 | 最后验证于: 2024 年 11 月 5 日

警告

本教程自 PyTorch 2.4 起已弃用。有关使用 Custom Operators 扩展 PyTorch 的最新指南，请参阅PyTorch Custom Operators。

Dispatcher 是 PyTorch 的内部组件，负责确定当你调用 torch::add 等函数时实际应该运行哪些代码。这可能很复杂，因为 PyTorch 操作需要处理许多“分层”在彼此之上的横切关注点。以下是一些它处理的事项示例

根据输入张量的设备，在 operator 的 CPU 和 CUDA 实现之间切换。
根据是否需要处理 autograd，在 operator 的 autograd 和后端实现之间切换。
在自动混合精度需要时应用 autocasting。
当 operator 在 vmap 调用下运行时，应用批处理规则。
如果你正在追踪模型以进行导出，则追踪操作的执行。

如果在你的自定义 operator 代码中，你发现自己手动编写 if 语句来处理这些情况，那么 dispatcher API 可以帮助组织你的代码。（反之，如果你的自定义 operator 非常简单，仅用于 CPU 推理，你可能不需要使用 dispatcher，只需使用基本 API 即可。）

在本教程中，我们将介绍如何构建自定义 operator 注册以使用 dispatcher 来组织各种组件。我们将假设你熟悉如何注册 operator 以及如何编写自定义 autograd 函数。

定义 schema 和后端实现¶

Dispatcher 的一般原理是它将 operator 的实现划分为多个 kernel，每个 kernel 实现特定 dispatch key 的功能，例如 CPU、CUDA。在你调用 operator 时，dispatcher 会确定最高优先级的 dispatch key 是什么（这是通过查看张量参数以及一些线程局部状态来完成的），并将控制权转移给该 dispatch key 的 kernel。最终效果是，当你调用 operator 时，我们首先执行 Autograd kernel，然后根据传入张量的设备类型重新调度到后端 kernel。

让我们看看实现这一过程涉及的各个部分。首先，我们必须定义相关 operator 的 schema。与简单的 pybind11 式 operator 注册不同，此时我们实际上并不提供 operator 的实现；我们只提供一个 schema 字符串，指定 operator 的类型签名，所有其他 kernel 都将遵循该签名

TORCH_LIBRARY(myops, m) {
  m.def("myadd(Tensor self, Tensor other) -> Tensor");
}

接下来，我们需要实际提供这个 operator 的一些实现。具体来说，这里有一个非常简单的 CPU 加法实现

Tensor myadd_cpu(const Tensor& self_, const Tensor& other_) {
  TORCH_CHECK(self_.sizes() == other_.sizes());
  TORCH_INTERNAL_ASSERT(self_.device().type() == DeviceType::CPU);
  TORCH_INTERNAL_ASSERT(other_.device().type() == DeviceType::CPU);
  Tensor self = self_.contiguous();
  Tensor other = other_.contiguous();
  Tensor result = torch::empty(self.sizes(), self.options());
  const float* self_ptr = self.data_ptr<float>();
  const float* other_ptr = other.data_ptr<float>();
  float* result_ptr = result.data_ptr<float>();
  for (int64_t i = 0; i < result.numel(); i++) {
    result_ptr[i] = self_ptr[i] + other_ptr[i];
  }
  return result;
}

我们希望将此函数注册为 myops::myadd 的一个实现。然而，简单的注册方法 (def("myadd", myadd_cpu)) 会在所有情况下注册该 kernel，即使张量不是 CPU 张量！(在内部，我们称这些为“catch-all” kernels，因为它们捕获所有情况。) 为了确保 myadd_cpu 仅对 CPU 张量运行，我们可以使用 TORCH_LIBRARY_IMPL 宏

TORCH_LIBRARY_IMPL(myops, CPU, m) {
  m.impl("myadd", myadd_cpu);
}

TORCH_LIBRARY_IMPL 允许我们在特定的 dispatch key (在此例中为 CPU) 上注册 operator 的实现。对 impl 的每次调用都会将一个 CPU kernel 与相应的 operator 相关联 (我们之前在 TORCH_LIBRARY 块中定义了该 operator)。如果我们也有一个 CUDA 实现 myadd_cuda，我们可以在另一个 TORCH_LIBRARY_IMPL 块中注册它

TORCH_LIBRARY_IMPL(myops, CUDA, m) {
  m.impl("myadd", myadd_cuda);
}

这些注册可以跨文件甚至跨库边界进行；例如，你可以将这两个 TORCH_LIBRARY_IMPL 块编译到单独的 myops_cpu 和 myops_cuda 动态库中。一般来说，你的注册结构将是这样的

一个单独的 TORCH_LIBRARY，在一个中心位置列出你的命名空间中的每个自定义 operator。
每个 dispatch key 对应一个 TORCH_LIBRARY_IMPL，用于注册该 key 的实现 (例如，CPU 或 CUDA)。如果你愿意，可以进一步将 TORCH_LIBRARY_IMPL 块细分为每个 operator 一个块。如果你每个 operator 实现都有一个单独的文件，但不想在头文件中公开 operator，这种方式很方便；你只需将注册放在定义 operator 的 cpp 文件中即可。

注意

你知道你也可以为 PyTorch 中现有的核心 operator 编写 TORCH_LIBRARY_IMPL 块吗？PyTorch 的 XLA 支持就是这样实现的：torch_xla 库包含一个 TORCH_LIBRARY_IMPL，它为 XLA dispatch key 上的所有基本 operator 提供实现。

对于不需要 autograd 的 operator¶

注意：本节仅适用于 PyTorch >= 1.10 版本。

在下一节中，我们将讨论如何为一个 operator 添加 autograd 支持。但对于不需要 autograd 支持的操作，应注册以下 kernel 以提高可用性，并使你的操作 behave like PyTorch 的内置 operator。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}

上面几行注册了一个 Autograd kernel，它在正向传播时附加一个哑 NotImplemented 节点（保留输入的 require_grad 属性）。在反向传播时，NotImplemented 节点会引发错误。这对于调试大型模型很有帮助，因为之前很难精确定位正向传播期间 requires_grad 属性丢失的位置。

In-place 或 view 操作¶

为确保正确性和最佳性能，如果你的操作会就地改变输入或返回与某个输入张量别名的张量，应采取以下两个额外步骤

除了上面的 Autograd kernel 外，还应注册一个 ADInplaceOrView kernel。此 kernel 处理必要的簿记工作，以确保就地或 view 操作的正确性。需要注意的是，此 ADInplaceOrView kernel 应仅与 autogradNotImplementedFallback 一起使用。

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl(op, autogradNotImplementedFallback());
}
TORCH_LIBRARY_IMPL(myops, ADInplaceOrView, m) {
  m.impl(op, autogradNotImplementedInplaceOrViewFallback());
}

上面注册的 Autograd 或 ADInplaceOrView boxed kernels 在其逻辑中依赖于 operator schema 信息。如果你的操作会就地改变输入或返回与某个输入张量别名的张量，务必确保你的 schema 正确反映了这一点。有关如何标注 schema 的更多信息，请参阅此处。

添加 autograd 支持¶

至此，我们有了一个同时包含 CPU 和 CUDA 实现的 operator。如何为其添加 autograd 支持呢？正如你可能猜到的，我们将注册一个 autograd kernel (类似于自定义 autograd 函数教程中描述的内容)！然而，这里有一个转折：与 CPU 和 CUDA kernel 不同，autograd kernel 需要重新调度 (redispatch)：它需要回调到 dispatcher 以获取推理 kernel，例如 CPU 或 CUDA 实现。

因此，在我们编写 autograd kernel 之前，让我们编写一个调度函数 (dispatching function)，它会调用 dispatcher 来查找你的 operator 的正确 kernel。此函数构成了你的 operator 的公共 C++ API——实际上，PyTorch C++ API 中的所有张量函数都在底层以相同的方式调用 dispatcher。下面是调度函数的示例

Tensor myadd(const Tensor& self, const Tensor& other) {
  static auto op = torch::Dispatcher::singleton()
    .findSchemaOrThrow("myops::myadd", "")
    .typed<decltype(myadd)>();
  return op.call(self, other);
}

让我们分解一下

在第一行中，我们从 dispatcher 中查找与我们要调度的 operator 对应的 typed operator handle。findSchemaOrThrow 接受两个参数：operator 的 (命名空间限定的) 名称，以及 operator 的重载名称 (通常是空字符串)。typed 将动态类型 handle 转换为静态类型 handle (通过运行时测试以确保你提供了正确的 C++ 类型)，这样我们就可以对其进行正常的 C++ 调用。我们传入 decltype(myadd)，因为调度函数的类型与注册到 dispatcher 的底层 kernel 的类型相同。

为了性能，此计算在静态变量中完成，这样我们只需进行一次（慢速）查找。如果你键入了错误的 operator 名称，第一次调用此函数时将出错。
在第二行中，我们只需使用传递给调度函数的所有参数 call (调用) operator handle。这实际上会调用 dispatcher，最终控制权将转移到适合此调用的任何 kernel。

有了调度函数，我们现在可以编写 autograd kernel 了

class MyAddFunction : public torch::autograd::Function<MyAddFunction> {
 public:
  static Tensor forward(
      AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {
    at::AutoNonVariableTypeMode g;
    return myadd(self, other);
  }

  static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) {
    auto grad_output = grad_outputs[0];
    return {grad_output, grad_output};
  }
};

Tensor myadd_autograd(const Tensor& self, const Tensor& other) {
  return MyAddFunction::apply(self, other)[0];
}

autograd 函数使用 torch::autograd::Function 正常编写，只是我们不是直接在 forward() 中编写实现，而是

使用 at::AutoNonVariableTypeMode RAII guard 关闭 autograd 处理，然后
调用调度函数 myadd 以回调到 dispatcher。

如果没有 (1)，你的调用将进入无限循环 (并导致栈溢出)，因为 myadd 会将你送回此函数 (因为最高优先级的 dispatch key 仍然是 autograd)。有了 (1)，autograd 将从考虑的 dispatch key 集合中排除，我们将转到下一个 handler，它将是 CPU 或 CUDA。

我们现在可以像注册 CPU/CUDA 函数一样注册此函数了

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
  m.impl("myadd", myadd_autograd);
}

注意

在此示例中，我们将 kernel 注册到 Autograd，这会将它安装为所有后端的 autograd kernel。你还可以通过使用相应的特定后端 dispatch key (例如，AutogradCPU 或 AutogradCUDA`) 来注册针对特定后端优化的 kernel。要更详细地探索这些以及其他 dispatch key 选项，请查阅 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。



超越 autograd¶
从某种意义上说，dispatcher 并没有做太多事情：它所做的只是实现了一个更高级的 if 语句，类似于这样
class MyAddFunction : ... {
public:
  static Tensor forward(
    AutogradContext *ctx, torch::Tensor self, torch::Tensor other) {

    if (self.device().type() == DeviceType::CPU) {
      return add_cpu(self, other);
    } else if (self.device().type() == DeviceType::CUDA) {
      return add_cuda(self, other);
    } else {
      TORCH_CHECK(0, "Unsupported device ", self.device().type());
    }
  }
  ...
}


那么为什么要使用 dispatcher 呢？有几个原因

它是去中心化的。你可以组装一个 operator 的所有部分（CPU、CUDA、Autograd），而无需编写一个引用所有部分的单一、集中的 if 语句。重要的是，第三方可以为其他方面注册额外的实现，而无需修改 operator 的原始定义。我们将在在 C++ 中为新后端扩展 dispatcher 中更详细地讨论扩展 dispatcher。
它支持比 CPU、CUDA 和 Autograd 更多的 dispatch key。你可以在 c10/core/DispatchKey.h 中看到 PyTorch 中当前实现的所有 dispatch key 列表。这些 dispatch key 为 operator 实现了各种可选功能，如果你决定希望你的自定义 operator 支持这些功能，你只需为相应的 key 注册一个 kernel。
dispatcher 实现了对 boxed fallback 函数的支持，这些函数可以实现一次并应用于系统中的所有 operator。Boxed fallback 可以用于为 dispatch key 提供默认行为；如果你使用 dispatcher 实现你的 operator，你也选择了所有这些操作的 fallback。

以下是一些你可能需要为其定义 operator 的特定 dispatch key。

Autocast¶
Autocast dispatch key 实现了对自动混合精度 (AMP) 的支持。一个 autocast wrapper kernel 通常会在运行操作之前将输入的 float16 或 float32 CUDA 张量转换为某种首选精度。例如，浮点 CUDA 张量上的 matmul 和卷积通常在 float16 中运行更快，使用更少的内存，同时不影响收敛。Autocast wrapper 仅在启用了 autocast 的上下文中生效。
这是一个假设的自定义 matmul 的 autocast wrapper，以及它的注册
// Autocast-specific helper functions
#include <ATen/autocast_mode.h>

Tensor mymatmul_autocast(const Tensor& self, const Tensor& other) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  return mymatmul(at::autocast::cached_cast(at::kHalf, self),
                  at::autocast::cached_cast(at::kHalf, other));
}

TORCH_LIBRARY_IMPL(myops, Autocast, m) {
  m.impl("mymatmul", mymatmul_autocast);
}


cached_cast(kHalf, tensor) 将 tensor 转换为 float16，如果 tensor 是 CUDA 且为 float32，否则保持 tensor 不变 (参见原生 autocasted 操作的资格策略)。这确保了如果网络在任何混合的 float16 和 float32 CUDA 张量上调用 mymatmul，mymatmul 将在 float16 中运行。同时，使用非 CUDA、整数类型或 float64 输入调用 mymatmul 不受影响。建议在自己的 autocast wrapper 中使用 cached_cast 遵循原生资格策略，但并非强制要求。例如，如果你想强制所有输入类型都以 float16 执行，可以使用 return mymatmul(self.half(), other.half()); 而非 cached_cast。
注意，就像我们的 autograd kernel 一样，我们在重新调度之前，从调度中排除 Autocast key。
默认情况下，如果没有提供 autocast wrapper，我们会直接 fallthrough 到常规 operator 实现（不会发生 autocasting）。（我们没有使用 myadd 作为此示例，因为逐点加法不需要 autocasting，只需 fall through 即可。）
何时应该注册 autocast wrapper？遗憾的是，对于操作的首选精度没有明确的规定。你可以通过查看cast lists 来了解一些原生操作的首选精度。一般指导意见

执行规约的操作应该可能以 float32 执行，
任何底层执行卷积或 gemm 的操作应该可能以 float16 执行，并且
其他具有多个浮点张量输入的操作应将其标准化为共同的精度 (除非实现支持不同精度的输入)。

如果你的自定义操作属于第三类，promote_type 模板有助于确定输入张量中存在的范围最广的浮点类型，这是执行类型的最安全选择
#include <ATen/autocast_mode.h>

Tensor my_multiple_input_op_autocast(const Tensor& t0, const Tensor& t1) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  // The required at::kHalf argument is an optimistic initial guess.
  auto exec_type = at::autocast::promote_type(at::kHalf, t0, t1);
  return my_multiple_input_op(at::autocast::cached_cast(exec_type, t0),
                              at::autocast::cached_cast(exec_type, t1));
}


如果你的自定义操作是启用 autograd 的，你只需要为注册 autograd wrapper 的相同名称编写并注册一个 autocast wrapper。例如，如果你想为 autograd 部分中显示的 myadd 函数创建一个 autocast wrapper，你只需要
Tensor myadd_autocast(const Tensor& self, const Tensor& other) {
  c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
  return myadd(at::autocast::cached_cast(<desired dtype>, self),
               at::autocast::cached_cast(<desired dtype>, other));
}

TORCH_LIBRARY_IMPL(myops, Autocast, m) {
  m.impl("myadd", myadd_autocast);
}


无需单独的复杂操作来使 backward 方法兼容 autocast。然而，自定义 autograd 函数中定义的 backward 方法将与 autocast 为 forward 方法设置的 dtype 相同，因此你应该选择一个适合你的 forward 和 backward 方法的 <desired dtype>。


批处理 (Batched)¶
批处理张量允许你以逐样本的方式编写代码，然后在 vmap 调用下运行时自动进行批处理。编写批处理规则的 API 目前正在开发中，但一旦稳定，你就可以通过在 Batched dispatch key 上注册 kernel 来为你的 operator 添加 vmap 支持。


追踪器 (Tracer)¶
Tracer dispatch key 实现了在运行 torch.jit.trace 时将 operator 调用记录到追踪中的支持。我们打算提供一个 boxed fallback 来实现任意操作的追踪，参见 issue #41478 以跟踪进度。

在 C++ 中注册一个 Dispatched Operator¶

定义 schema 和后端实现¶

对于不需要 autograd 的 operator¶

In-place 或 view 操作¶

添加 autograd 支持¶

超越 autograd¶

Autocast¶

批处理 (Batched)¶

追踪器 (Tracer)¶

文档

教程

资源