PyTorch 2.2 中文官方教程（十三）（1）-阿里云开发者社区

在 C++中注册一个分发的运算符

原文：pytorch.org/tutorials/advanced/dispatcher.html

译者：飞龙

协议：CC BY-NC-SA 4.0

分发器是 PyTorch 的一个内部组件，负责确定在调用诸如torch::add这样的函数时实际运行哪些代码。这可能并不简单，因为 PyTorch 操作需要处理许多“层叠”在彼此之上的交叉关注点。以下是它处理的一些事项的示例：

根据输入张量的设备在 CPU 和 CUDA 实现之间切换运算符。
在是否需要自动微分处理的情况下，在运算符的自动微分和后端实现之间切换。
在需要自动混合精度时应用自动转换。
在vmap调用下运行运算符时应用批处理规则。
跟踪操作的执行，如果您正在跟踪一个模型以进行导出。

如果在您的自定义运算符代码中发现自己手动编写 if 语句来处理这些情况，分发器 API 可以帮助组织您的代码。（相反，如果您的自定义运算符非常简单且仅用于 CPU 推断，则可能不需要使用分发器，只需使用基本 API。）

在本教程中，我们将描述如何结构化自定义运算符注册以使用分发器来组织各种组件。我们假设您已经熟悉如何注册运算符以及如何编写自定义自动微分函数。

定义模式和后端实现

分发器背后的一般原则是将运算符的实现分成多个内核，每个内核为特定的分发键实现功能，例如 CPU、CUDA。分发器确定在调用运算符时最高优先级的分发键是什么（这是通过查看张量参数以及一些线程本地状态来完成的），并将控制权转移到该分发键的内核。最终效果是当您调用一个运算符时，我们首先执行自动微分内核，然后根据传入张量的设备类型重新分发到后端内核。

让我们看看使这一切发生所涉及的各个部分。首先，我们必须为所讨论的运算符定义模式。与简单的 pybind11 风格的运算符注册不同，我们此时实际上并没有提供我们运算符的实现；我们只提供一个模式字符串，指定所有其他内核将遵守的运算符类型签名：

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

接下来，我们需要实际提供一些这个运算符的实现。具体来说，这是一个在 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)）会注册内核在所有情况下运行，即使张量不是 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动态库中。一般来说，您的注册结构将如下所示：

一个单独的TORCH_LIBRARY，列出您命名空间中的每个自定义操作符，集中在一个地方。
每个调度键注册一个TORCH_LIBRARY_IMPL，为该键（例如，CPU 或 CUDA）注册实现。如果愿意，您还可以将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属性的位置。

原地或视图操作

为确保正确性和最佳性能，如果您的操作在原地更改输入或返回一个与输入之一别名的张量，则应采取两个额外步骤：

除了上面的Autograd内核外，还注册一个ADInplaceOrView内核。该内核处理必要的记录工作，以确保原地或视图操作的正确性。重要的是要注意，此 ADInplaceOrView 内核应仅与autogradNotImplementedFallback一起使用。

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

上面注册的Autograd或ADInplaceOrView封装的内核依赖于其逻辑中的运算符模式信息。如果您的操作在原地对输入进行了更改，或者返回一个与输入之一别名的张量，那么确保您的模式正确反映这一点非常重要。请参阅此处以获取有关如何注释模式的更多信息。

添加自动求导支持

到目前为止，我们有一个既有 CPU 实现又有 CUDA 实现的操作符。我们如何为其添加自动求导支持？正如您可能猜到的那样，我们将注册一个自动求导内核（类似于自定义自动求导函数教程中描述的内容）！但是，有一个转折：与 CPU 和 CUDA 内核不同，自动求导内核需要重新调度：它需要回调到调度程序以获取推断内核，例如 CPU 或 CUDA 实现。

因此，在编写自动求导内核之前，让我们编写一个调度函数，该函数调用调度程序以找到适合您操作符的正确内核。这个函数构成了您操作符的公共 C++ API - 实际上，PyTorch 的 C++ API 中的所有张量函数都在底层以相同的方式调用调度程序。调度函数如下所示：

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

让我们来详细了解一下：

在第一行中，我们从调度程序中查找与我们要分派的运算符对应的类型化运算符句柄。findSchemaOrThrow接受两个参数：运算符的（命名空间限定的）名称和运算符的重载名称（通常为空字符串）。typed将动态类型的句柄转换为静态类型的句柄（进行运行时测试以确保您提供了正确的 C++类型），以便我们可以对其进行正常的 C++调用。我们传递decltype(myadd)，因为分派函数的类型与注册到调度程序的基础内核的类型相同。
为了性能，此计算是在静态变量中完成的，因此我们只需要进行一次（慢速）查找。如果您拼错了要调用的运算符的名称，那么在第一次调用此函数时，此查找将出错。
在第二行中，我们简单地使用传递给分派函数的所有参数“调用”运算符句柄。这将实际调用调度程序，最终控制将转移到适用于此调用的任何内核。

有了分派函数，我们现在可以编写自动微分内核了：

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

自动微分函数是使用torch::autograd::Function正常编写的，只是在forward()中不直接编写实现，而是：

使用at::AutoNonVariableTypeMode RAII 保护关闭自动微分处理，然后
调用分派函数myadd以回调到调度程序。

没有（1），您的调用将无限循环（并堆栈溢出），因为myadd将将您发送回此函数（因为最高优先级的调度键仍然是自动微分）。有了（1），自动微分将从考虑的调度键集合中排除，我们将转到下一个处理程序，这将是 CPU 和 CUDA。

我们现在可以以与注册 CPU/CUDA 函数相同的方式注册此函数：

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

注意

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

超越自动微分

在某种意义上，调度程序并没有做太多事情：它只是实现了一个类似于这样的 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());
  }
  }
  ...
}

为什么要使用调度程序？有几个原因：

它是分散的。您可以组装运算符的所有部分（CPU、CUDA、Autograd）而无需编写一个引用所有这些部分的单个集中 if 语句。重要的是，第三方可以注册其他方面的额外实现，而无需修补运算符的原始定义。我们将在扩展调度程序以支持新后端中更多地讨论扩展调度程序。
它支持比 CPU、CUDA 和 Autograd 更多的调度键。您可以在 PyTorch 中当前实现的c10/core/DispatchKey.h中看到当前实现的所有调度键的完整列表。这些调度键为运算符实现了各种可选功能，如果您决定希望您的自定义运算符支持此功能，您只需为适当的键注册一个内核。
调度程序实现了对装箱回退函数的支持，这些函数可以一次实现并应用于系统中的所有运算符。装箱回退可用于为调度键提供默认行为；如果您使用调度程序来实现您的运算符，您还可以选择为所有这些操作启用回退。

以下是一些您可能需要为其定义运算符的特定调度键。

自动转换

Autocast 分派键实现了对自动混合精度（AMP）的支持。自动转换包装器内核通常会将传入的float16或float32 CUDA 张量转换为某种首选精度，然后运行操作。例如，在浮点 CUDA 张量上运行的矩阵乘法和卷积通常在float16中运行更快，使用更少的内存，而不会影响收敛。自动转换包装器仅在启用自动转换的上下文中起作用。

以下是一个假设的自定义矩阵乘法的自动转换包装器，以及其注册：

// 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保持不变（参见natively autocasted ops 的资格政策）。这确保了如果网络在任何混合float16和float32 CUDA 张量上调用mymatmul，mymatmul将以float16运行。同时，对于非 CUDA、整数类型或float64输入的mymatmul调用不受影响。建议在自己的自动转换包装器中使用cached_cast遵循本机资格政策，但不是必需的。例如，如果您想要强制所有输入类型执行float16，您可以使用return mymatmul(self.half(), other.half());而不是使用cached_cast。

请注意，与我们的自动求导内核一样，在重新分派之前，我们将Autocast键排除在分派之外。

默认情况下，如果没有提供自动转换包装器，我们将直接转到常规操作员实现（不会发生自动转换）。（我们没有在此示例中使用myadd，因为逐点加法不需要自动转换，应该直接通过。）

何时应注册自动转换包装器？不幸的是，没有关于操作首选精度的明确规则。您可以通过查看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));
}

如果您的自定义操作是自动求导启用的，您只需要为与自动求导包装器注册的相同名称编写并注册一个自动转换包装器。例如，如果您想要一个myadd函数的自动转换包装器，只需

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

没有单独的技巧使得反向方法与自动转换兼容。但是，您自定义的自动求导函数中定义的反向方法将以与自动转换为前向方法设置的相同 dtype 运行，因此您应该选择一个适合您的前向和反向方法的。

批处理

批处理张量允许您以每个示例的方式编写代码，然后在vmap调用下运行时自动批处理它们。编写批处理规则的 API 目前正在开发中，但一旦稳定下来，您可以通过在批处理分派键上注册一个内核来为您的操作添加对vmap的支持。

追踪器

追踪器分派键实现了在运行torch.jit.trace时将操作调用记录到跟踪中的支持。我们打算提供一个包装回退，用于实现任意操作的跟踪，参见issue＃41478以跟踪进展。

在 C++中为新后端扩展调度程序

原文：pytorch.org/tutorials/advanced/extend_dispatcher.html

译者：飞龙

协议：CC BY-NC-SA 4.0

在本教程中，我们将逐步介绍扩展调度程序的所有必要步骤，以添加一个位于pytorch/pytorch存储库之外的新设备，并保持与原生 PyTorch 设备同步。在这里，我们假设您熟悉如何在 C++中注册调度运算符以及如何编写自定义自动微分函数。

注意

本教程涉及 PyTorch 内部许多正在积极改进的组件，请在决定跟随本教程时预期 API 的更改。我们将保持本教程与最新的 API 保持同步。

什么是新后端？

向 PyTorch 添加一个新后端需要来自后端扩展者的大量开发和维护。在添加新后端之前，让我们首先考虑一些常见用例和推荐的解决方案：

如果您有现有 PyTorch 运算符的新算法，请向 PyTorch 发送一个 PR。
如果您想提出一个新的运算符，请向 PyTorch 发送一个功能请求/PR。
如果您想为新设备/硬件（如 Google TPU 和定制芯片）添加支持，通常需要使用特定于硬件的 API 来编写内核，请按照本教程并向 PyTorch 添加一个树外后端。
如果您想为现有运算符添加支持，但使用不同的张量布局/表示，如稀疏和量化，这将强制您的内核以更有效的方式编写，考虑到布局/表示限制，请按照本教程并向 PyTorch 添加一个树外后端。

在本教程中，我们将主要关注添加一个新的树外设备。为不同张量布局添加树外支持可能与设备共享许多常见步骤，但我们尚未看到这种集成的示例，因此可能需要 PyTorch 进行额外的工作来支持它。

为您的后端获取一个调度键

PyTorch 运算符是用 C++实现的，并通过 Python 绑定在 Python 前端中提供。PyTorch 调度程序将运算符的实现分为多个内核，每个内核与特定的调度键相关联。在 PyTorch 中支持一个新后端基本上意味着为 C++中的每个 PyTorch 运算符编写一个内核，然后将它们注册到调度程序中代表您定制后端的调度键。

调度键是调度系统中的标识符。调度程序查看输入张量上携带的调度键，并相应地调用正确的内核。PyTorch 为原型化树外后端扩展提供了三个预留的调度键（以及它们对应的 Autograd 键）：

PrivateUse1/AutogradPrivateUse1
PrivateUse2/AutogradPrivateUse2
PrivateUse3/AutogradPrivateUse3

您可以选择上述任何键来原型化您的定制后端。要在PrivateUse1后端上创建一个张量，您需要在TensorImpl构造函数中设置调度键。

/* Example TensorImpl constructor */
TensorImpl(
  Storage&&  storage,
  DispatchKeySet  ks,
  const  caffe2::TypeMeta  data_type);
// To create a TensorImpl on PrivateUse1 backend, pass in the following ks to TensorImpl creation.
DispatchKeySet  ks  =  c10::DispatchKeySet{c10::DispatchKey::PrivateUse1,  c10::DispatchKey::AutogradPrivateUse1};

请注意，上面的TensorImpl类假定您的张量由类似 CPU/CUDA 的存储支持。我们还提供了OpaqueTensorImpl，用于没有存储的后端。您可能需要调整/覆盖某些方法以适应您的定制硬件。PyTorch 存储库中的一个示例是Vulkan TensorImpl。

注意

一旦原型完成，并且您计划为您的后端扩展进行定期发布，请随时向pytorch/pytorch提交一个 PR，以保留一个专用的调度键给您的后端。

流畅的 Python 第二版（GPT 重译）（十一）（2）https://developer.aliyun.com/article/1482564

PyTorch 2.2 中文官方教程（十三）（1）

在 C++中注册一个分发的运算符

定义模式和后端实现

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

原地或视图操作

添加自动求导支持

超越自动微分

自动转换

批处理

追踪器

在 C++中为新后端扩展调度程序

什么是新后端？

为您的后端获取一个调度键

热门文章

最新文章

相关电子书

推荐镜像

探索云世界

热门

云计算

大数据

云原生

人工智能

数据库

开发与运维

活动广场

任务中心

训练营

直播

乘风者计划

下载

镜像站

技术资料

PyTorch 2.2 中文官方教程（十三）（1）

在 C++中注册一个分发的运算符

定义模式和后端实现

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

原地或视图操作

添加自动求导支持

超越自动微分

自动转换

批处理

追踪器

在 C++中为新后端扩展调度程序

什么是新后端？

为您的后端获取一个调度键

热门文章

最新文章

相关电子书

推荐镜像