OP Lowering 指南¶

PyTorch 封装了 C++ ATen 张量库，该库在 GPU 和 CPU 上提供了广泛的操作实现。Pytorch/XLA 是 PyTorch 的扩展；其目的之一是将 PyTorch 操作转换为 XLA 操作。Lowering 定义了将高级表示形式转换为低级表示形式的过程。在本文档中，我将把 PyTorch 操作转换为 XLA 操作的过程称为 lowering。XLA 编译器也会将 XlaOp lowering 为 HLO，但这超出了本文档的范围。我们将把尚未提供 XLA lowering 的操作转发到 CPU 并调用 ATen 实现。转发到 CPU 的操作将导致明显的减速。我们必须 lowering 模型中使用的所有操作，以实现最佳性能。

以下是您可能会从 PyTorch/XLA 调试工具中看到的尚未 lowering 的操作示例

pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss, aten::_ctc_loss_backward,  Please open a GitHub issue with the above op lowering requests.

开始之前¶

您应该遵循为 Pytorch/XLA 做贡献中的说明，安装所需的依赖项并从源代码构建 pytorch 和 pytorch/XLA。您无需访问 TPU 即可实现 lowering。建议在工作站上进行实验，并将其配置为使用 XLA:CPU。您可以通过运行以下命令将 Pytorch/XLA 配置为使用 XLA:CPU

export PJRT_DEVICE=CPU

理解操作¶

您可以在native_functions.yaml中找到 C++ ATen 操作的定义。从源代码构建 Pytorch/XLA 后，您还将在 xla/torch_xla/csrc/aten_fallback.h/cpp 中找到我们的默认实现（一个 boxed kernel，它将调用转发到 PyTorch 原生内核）。Pytorch 操作通常可以轻松映射到 PyTorch 张量 API。如果不是这种情况，建议在 PyTorch repo 下搜索 PyTorch 原生实现。目标是将 PyTorch 操作 lowering 为 XLA 操作语义中定义的一系列 XLA 操作。

文件结构¶

以下提到的所有文件都位于 xla/torch_xla/csrc 文件夹下，除了 codegen/xla_native_functions.yaml

xla_native_functions.yaml 包含显式 lowering 的所有运算符列表（来自 Core Aten list）。组合运算符未在此处列出。此处的每个运算符名称都必须直接匹配 native_functions.yaml 中列出的 pytorch 运算符。此文件充当添加新 xla 运算符的接口，并且是 PyTorch 的 codegen machinery 的输入。它生成以下 3 个文件：XLANativeFunctions.h、RegisterXLA.cpp 和 RegisterAutogradXLA.cpp
XLANativeFunctions.h 和 aten_xla_type.cpp 是 PyTorch 进入 pytorch_xla 世界的入口点，并包含每个运算符的手动编写的 XLA lowering。 XLANativeFunctions.h 是通过 xla_native_functions.yaml 和 PyTorch 核心 native_functions.yaml 文件的组合自动生成的，并包含需要在 aten_xla_type.cpp 中定义的内核的声明。此处编写的内核需要使用输入 at::Tensor 和其他参数构造 ‘XLATensor’。生成的 XLATensor 需要在返回到 PyTorch 世界之前转换回 at::Tensor。
RegisterXLA.cpp 和 RegisterAutogradXLA.cpp 是自动生成的文件，用于将所有 lowering 注册到 PyTorch Dispatcher。它们还包括 out= 和 inplace 运算符的自动生成的包装器实现。
aten_fallback.h/.cpp 包含我们的 boxed fallback 实现。如果未在 xla_native_functions.yaml + aten_xla_type.cpp 中显式定义 lowering，并且运算符不是复合运算符，则将使用 boxed fallback 内核。
tensor_methods.h 包含 XLATensor 声明。这些声明通常与我们在 XLANativeFunctions.h 中声明的 at::Tensor 节点一一对应
tensor_methods.cpp 包含 tensor_methods.h 中定义的 XLATensor node 的实现。我们从参数的 ir::Value 构造了相应的 ir::op，并将其包装在 XLATensor 中。Ir 代表中间表示。
ops/ 目录包含所有 ir::ops 声明和定义。较小的节点可以放在 ops/ops.h/.cpp 中。更复杂的节点可以放在单独的文件中。所有 ops 都继承自 ir::ops::Node，并提供一种将输入 ir::Value lowering 为 XlaOp 序列的方法。

单元测试¶

我们的 CI 每天为每个更改运行 PyTorch 原生 python 测试。如果我们提供 lowering，这些测试将使用 XLA 实现。通常我们不需要为 PyTorch/XLA 添加额外的 python 测试，除非我们想验证一些 xla 行为（如动态形状）或者由于某种原因跳过了 pytorch 原生测试。如果需要，应将 python 测试添加到 xla/test/test_operations.py 中。我们还需要在 xla/test/cpp/test_aten_xla_tensor.cpp 中添加 CPP 测试。此测试应调用 PyTorch c++ API，并验证我们的实现是否产生与 PyTorch 原生实现相同的结果。我们还需要通过检查 aten::op 和 xla::op 计数器来验证当张量是 XLA 张量时是否调用了 xla 实现。

提示¶

lowering 的过程是将 PyTorch 操作分解为 XlaOp 序列。为了提供 PyTorch 操作的良好 lowering，需要很好地掌握 XLA 的功能。阅读 XlaOp 文档并查看类似 ops 的 lowering 方式是实现这一目标的最佳方法。您可以在此 Op lowering PR中找到一个最小的 Op lowering 示例。您还可以在此 backward lowering PR中找到一个稍微复杂的带有向后 lowering 的示例。

我们在 RegisterXLA.cpp 中为某些运算符自动生成了 out= 和 inplace 运算符的包装器实现。在这种情况下，我们只需要 lowering vanilla op。一个例子是 lerp 运算符，它在 native_functions.yaml 中有 6 个变体，它们是

- lerp_.Scalar
- lerp_.Tensor
- lerp.Scalar_out
- lerp.Tensor_out
- lerp.Scalar
- lerp.Tensor

并将生成函数原型

at::Tensor lerp(const at::Tensor & self, const at::Tensor & end, const at::Scalar & weight);
at::Tensor & lerp_(at::Tensor & self, const at::Tensor & end, const at::Scalar & weight);
at::Tensor lerp(const at::Tensor & self, const at::Tensor & end, const at::Tensor & weight);
at::Tensor & lerp_out(const at::Tensor & self, const at::Tensor & end, const at::Tensor & weight, at::Tensor & out);
at::Tensor & lerp_(at::Tensor & self, const at::Tensor & end, const at::Tensor & weight);
at::Tensor & lerp_out(const at::Tensor & self, const at::Tensor & end, const at::Scalar & weight, at::Tensor & out);

在 XLANativeFunctions.h 中，如果我们将它们全部添加到 xla_native_functions.yaml 中。但是，如果我们仅 lowering lerp.Scalar 和 lerp.Tensor 并检查 RegisterXLA.cpp，我们将看到

namespace {

at::Tensor wrapper_Scalar_lerp(const at::Tensor & self, const at::Tensor & end, const at::Scalar & weight) {
    // No device check


  // DeviceGuard omitted
  return torch_xla::lerp(self, end, weight);
}

} // anonymous namespace

at::Tensor & wrapper_Scalar_lerp_(at::Tensor & self, const at::Tensor & end, const at::Scalar & weight) {
  auto wrapper_Scalar_lerp__tmp = wrapper_Scalar_lerp(self, end, weight);
  at::_copy_from(wrapper_Scalar_lerp__tmp, self);
  return self;
}

...
  m.impl("lerp_.Scalar",
  TORCH_FN(wrapper_Scalar_lerp_));

codegen 将自动为使用我们的 lerp.Scalar 实现的 lerp_.Scalar 和 lerp.Scalar_out 生成 lowering，而无需我们提供显式 lowering。

一般来说，如果 pytorch 核心中有一个运算符既有 out-of-place 变体又有 out= 变体，那么最好为 out-of-place 变体编写 lowering，因为您将免费获得代码生成的 out= lowering。

对于每个节点，我们需要传递一个 ir::OpKind。这是一个 (示例)。您可以在 interned_strings.h 中找到 OpKind 定义。如果 aten 符号丢失，您可以提交像这样的 PR。

OP Lowering 指南¶

开始之前¶

理解操作¶

文件结构¶

单元测试¶

提示¶

文档

教程

资源