在 C++ 中注册一个分派操作符

定义模式和后端实现

调度器背后的基本原理是将操作符的实现划分为多个内核，每个内核负责实现特定调度键的功能，例如 CPU、CUDA。调度器会在调用操作符时确定最高优先级的调度键（这是通过查看张量参数以及一些线程局部状态来完成的），然后将控制权转移到该调度键对应的内核。最终的效果是，当您调用一个操作符时，我们会先执行 Autograd 内核，然后根据传入张量的设备类型重新分派到后端内核。

让我们来看一下实现这一目标的各个部分。首先，我们必须为所讨论的操作符定义模式。与简单的 pybind11 风格操作符注册不同，我们实际上并不在这一步骤提供操作符的实现；我们只是提供一个模式字符串，指定操作符的类型签名，所有其他内核都将遵循该签名：

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

接下来，我们需要实际提供这个操作符的一些实现。为了具体化，这里是一个非常简单的 CPU 上的加法实现：

Tensormyadd_cpu(constTensor&self_,constTensor&other_){
TORCH_CHECK(self_.sizes()==other_.sizes());
TORCH_INTERNAL_ASSERT(self_.device().type()==DeviceType::CPU);
TORCH_INTERNAL_ASSERT(other_.device().type()==DeviceType::CPU);
Tensorself=self_.contiguous();
Tensorother=other_.contiguous();
Tensorresult=torch::empty(self.sizes(),self.options());
constfloat*self_ptr=self.data_ptr<float>();
constfloat*other_ptr=other.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();
for(int64_ti=0;i<result.numel();i++){
result_ptr[i]=self_ptr[i]+other_ptr[i];
}
returnresult;
}

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

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

TORCH_LIBRARY_IMPL 允许我们为特定调度键（在本例中为 CPU）上的操作符注册实现。每次调用 impl 都会将一个 CPU 内核与相应的操作符关联起来（我们之前在 TORCH_LIBRARY 块中定义了该操作符）。如果我们还有一个 CUDA 实现 myadd_cuda，我们可以在一个单独的 TORCH_LIBRARY_IMPL 块中注册它：

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

这些注册可以分散在不同的文件中，甚至可以跨库边界；例如，您可以将这两个 TORCH_LIBRARY_IMPL 块编译到单独的 myops_cpu 和 myops_cuda 动态库中。一般来说，注册的结构将如下所示：

1. 一个统一的 TORCH_LIBRARY，用于在集中位置列出您命名空间中的所有自定义操作符。

2. 每个调度键（如 CPU 或 CUDA）对应一个 TORCH_LIBRARY_IMPL，用于为该键注册实现。如果您愿意，可以进一步将 TORCH_LIBRARY_IMPL 块细分为每个操作符的块。如果您为每个操作符实现使用单独的文件，但不想在头文件中暴露这些操作符，这会非常方便；您只需将注册放在定义操作符的 cpp 文件中即可。

您知道吗？您也可以为 PyTorch 中现有的核心操作符编写 TORCH_LIBRARY_IMPL 块。这就是 PyTorch 中 XLA 支持的实现方式：torch_xla 库包含一个 TORCH_LIBRARY_IMPL，它为 XLA 调度键上的所有基本操作符提供了实现。

对于不需要自动求导的操作符

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

在下一节中，我们将讨论如何为操作符添加自动求导支持。但对于不需要自动求导支持的操作，应注册以下内核以提高可用性，并使您的操作符行为类似于 PyTorch 的内置操作符。

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

以上代码注册了一个 Autograd 内核，该内核在前向传播时附加一个虚拟的 NotImplemented 节点（保留输入的 require_grad 属性）。在反向传播时，NotImplemented 节点会引发一个错误。这对于在较大的模型中进行调试很有帮助，因为在之前可能很难精确定位在前向传播过程中 requires_grad 属性是在哪里丢失的。

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

添加自动求导支持

此时，我们已经有了一个包含 CPU 和 CUDA 实现的算子。如何为其添加自动求导支持呢？正如您可能猜到的，我们将注册一个自动求核（类似于自定义自动求导函数教程中描述的内容）！不过，这里有一个区别：与 CPU 和 CUDA 核不同，自动求核需要重新分发：它需要回调到分发器以获取推理核，例如 CPU 或 CUDA 实现。

因此，在编写自动求核之前，我们先编写一个分发函数，该函数调用分发器以找到适合您算子的核。这个函数构成了您算子的公共 C++ API——事实上，PyTorch C++ API 中的所有张量函数在底层都以相同的方式调用分发器。以下是分发函数的样子：

Tensormyadd(constTensor&self,constTensor&other){
staticautoop=torch::Dispatcher::singleton()
.findSchemaOrThrow("myops::myadd","")
.typed<decltype(myadd)>();
returnop.call(self,other);
}

让我们来详细分解一下：

• 在第一行中，我们从调度器中查找一个类型化的操作符句柄，该句柄对应我们将要调度到的操作符。findSchemaOrThrow 接受两个参数：操作符的（命名空间限定的）名称和操作符的重载名称（通常为空字符串）。typed 将动态类型的句柄转换为静态类型的句柄（并进行运行时测试以确保您提供了正确的 C++ 类型），以便我们可以对其进行正常的 C++ 调用。我们传递 decltype(myadd)，因为调度函数的类型与注册到调度器的底层内核的类型相同。

  出于性能考虑，此计算在静态变量中完成，因此我们只需执行一次（较慢的）查找。如果您拼错了要调用的操作符名称，则在首次调用此函数时，此查找将出错。

• 在第二行中，我们简单地使用传递给调度函数的所有参数 call 操作符句柄。这将实际调用调度器，最终控制权将转移到适合此调用的内核。

有了 dispatch 函数，我们现在可以编写自动求导内核了：

classMyAddFunction:publictorch::autograd::Function<MyAddFunction>{
public:
staticTensorforward(
AutogradContext*ctx,torch::Tensorself,torch::Tensorother){
at::AutoNonVariableTypeModeg;
returnmyadd(self,other);
}

statictensor_listbackward(AutogradContext*ctx,tensor_listgrad_outputs){
autograd_output=grad_outputs[0];
return{grad_output,grad_output};
}
};

Tensormyadd_autograd(constTensor&self,constTensor&other){
returnMyAddFunction::apply(self,other)[0];
}

自动求导函数的编写与普通的 torch::autograd::Function 类似，区别在于我们不会直接在 forward() 中编写实现，而是：

1. 使用 at::AutoNonVariableTypeMode RAII 防护关闭自动求导处理，然后

2. 调用分发函数 myadd 以回退到分发器。

如果没有(1)，您的调用将会无限循环（并导致栈溢出），因为 myadd 会将您送回此函数（因为最高优先级的调度键仍然是 autograd）。有了(1)，autograd 将从考虑的调度键集合中排除，我们将转到下一个处理程序，即 CPU 或 CUDA。

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

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

在这个示例中，我们将内核注册到 Autograd，这会将其安装为所有后端的自动求导内核。您还可以通过使用相应的后端特定的调度键来为特定后端注册优化的内核，例如 AutogradCPU 或 AutogradCUDA。要更详细地了解这些及其他调度键选项，请查看 torch/_python_dispatcher.py 中提供的 PythonDispatcher 工具。

超越自动求导

从某种意义上说，调度器并没有做太多事情：它所做的只是实现了一个增强版的 if 语句，类似于这样：

classMyAddFunction:...{
public:
staticTensorforward(
AutogradContext*ctx,torch::Tensorself,torch::Tensorother){

if(self.device().type()==DeviceType::CPU){
returnadd_cpu(self,other);
}elseif(self.device().type()==DeviceType::CUDA){
returnadd_cuda(self,other);
}else{
TORCH_CHECK(0,"Unsupported device ",self.device().type());
}
}
...
}

那么为什么要使用 dispatcher 呢？有以下几点原因：

1. 它是去中心化的。您可以在不编写一个集中式的、引用所有部分（CPU、CUDA、Autograd）的 if 语句的情况下，组装操作符的所有部分。重要的是，第三方可以为其他方面注册额外的实现，而无需修改操作符的原始定义。我们将在为新后端扩展调度器中进一步讨论如何扩展调度器。

2. 它支持比 CPU、CUDA 和 Autograd 更多的调度键。您可以在 c10/core/DispatchKey.h 中查看当前在 PyTorch 中实现的所有调度键的完整列表。这些调度键为操作符实现了各种可选功能，如果您决定让自定义操作符支持这些功能，只需为适当的键注册一个内核即可。

3. 调度器实现了对盒式回退函数的支持，这些函数可以一次性实现并应用于系统中的所有操作符。盒式回退可用于为调度键提供默认行为；如果您使用调度器来实现操作符，您也将自动获得所有这些操作的回退功能。

以下是一些您可能需要为其定义操作符的特殊分发键。

// Autocast-specific helper functions
#include<ATen/autocast_mode.h>

Tensormymatmul_autocast(constTensor&self,constTensor&other){
c10::impl::ExcludeDispatchKeyGuardno_autocast(c10::DispatchKey::Autocast);
returnmymatmul(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);
}

#include<ATen/autocast_mode.h>

Tensormy_multiple_input_op_autocast(constTensor&t0,constTensor&t1){
c10::impl::ExcludeDispatchKeyGuardno_autocast(c10::DispatchKey::Autocast);
// The required at::kHalf argument is an optimistic initial guess.
autoexec_type=at::autocast::promote_type(at::kHalf,t0,t1);
returnmy_multiple_input_op(at::autocast::cached_cast(exec_type,t0),
at::autocast::cached_cast(exec_type,t1));
}

Tensormyadd_autocast(constTensor&self,constTensor&other){
c10::impl::ExcludeDispatchKeyGuardno_autocast(c10::DispatchKey::Autocast);
returnmyadd(at::autocast::cached_cast(<desireddtype>,self),
at::autocast::cached_cast(<desireddtype>,other));
}

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