目录
  • 文档 >
  • Codegen 迁移指南
快捷方式

Codegen 迁移指南

随着 PyTorch/XLA 迁移到 LTC (Lazy Tensor Core)，我们需要清理用于实现 op 降级的现有存根代码（分布在 6 个以上的文件中）。旧 op 降级的完整流程和文件结构可以在 op 降级指南 :ref:’op-lowering’ 中找到。

有关不同操作配置的更多信息，请参阅 codegen/xla_native_functions.yaml。将受支持的操作（在 supported 配置下）替换为 codegen 等效项（在 full_codegen 配置下）**不应**引入任何新的行为，这纯粹是为了清理目的。其他配置下的操作可能具有不同的行为。有关其他配置的信息，请参阅 codegen/xla_native_functions.yaml。

有关 PyTorch 中 dispatching 的更多信息，请参阅 Exyang 的博客文章。

开始之前

您应该按照 此处 的说明安装所需的依赖项并从源构建 pytorch 和 pytorch/XLA。您不需要 TPU 访问权限即可实现降级。建议在工作站上进行实验，并将其配置为使用 XLA:CPU。您可以通过运行以下命令将 Pytorch/XLA 配置为使用 XLA:CPU：

export PJRT_DEVICE=CPU

还建议您在进行 codegen 工作之前熟悉我们的 op 降级过程。

PyTorch/XLA 使用 https://github.com/pytorch/xla/issues/3560 来跟踪 codegen 迁移的状态。在进行 codegen 工作时，请在 issue 中附上您的 GitHub 别名和 PR 链接，以避免重复工作。

文件结构

下面提到的所有文件都位于 xla/torch_xla/csrc 文件夹下，xla_native_functions.yaml 除外。

PyTorch Codegen 文件

  • torch/csrc/lazy/core/shape_inference.h

    • 为每个 op 定义的形状推理函数，它接受 torch::lazy::shapes 作为输入并返回 torch::lazy::shape 作为输出。只有非结构化 op 才需要手动形状推理函数。

  • torchgen/gen_lazy_tensor.py

    • 构建在所有 ATen 后端使用的数据模型和助手之上，并为 lazy tensor 后端添加新功能。run_gen_lazy_tensor 在此文件中定义。

  • torchgen/dest/lazy_ir.py

    • 包含可由后端覆盖的数据类 GenLazyIR，并定义生成的 IR 类。

PyTorch/XLA Codegen 文件

  • xla/xla_native_functions.yaml

    • 包含 XLA 今天支持的所有 op。大多数 op 都属于 supported 类别，本文档的目的是将大多数 op 移至 full_codegen 类别。

  • xla/scripts/gen_lazy_tensor.py

    • 提供 codegen Codegen 类必需的 XLA 版本，并调用上游 codegen API。

  • xla/torch_xla/csrc/XLANativeFunctions.cpp

    • xla/codegen/xla_native_functions.yaml 的 full_codegen 列的结果。此处定义的 op 函数将实现 XLANativeFunctions.h 中声明的 op。每个 op 都将接收 at::tensor 并返回另一个包装在 XLATensor 中的 at::tensor。请注意，这部分以前是在 tensor_method.cpp 中手动完成的。

  • xla/torch_xla/csrc/LazyIr.h

    • xla/codegen/xla_native_functions.yaml 的 full_codegen 列的结果。定义用于构建 full_codegen op 的 IR。

PyTorch/XLA 旧 op 降级文件

  • xla/torch_xla/csrc/generated/aten_xla_type.cpp

    • 手动实现 xla/codegen/xla_native_functions.yaml 中定义的 op。将被 XLANativeFunctions.cpp 替换。

  • xla/torch_xla/csrc/generated/tensor.h

    • 定义 XLATensor 类和 XLATensor 方法声明。这些声明通常是 XLANativeFunctions.h 中声明的 at::Tensor 节点的逐一映射。对于 full_codegen op，将删除 XLATensor 方法。

  • xla/torch_xla/csrc/generated/tensor_method.cpp

    • 实现 tensor.h 中定义的张量方法。对于 full_codegen op，此文件将被删除。

  • xla/torch_xla/csrc/generated/ops/…

    • 定义“大多数” op 的 IR 类。可能多个 op 共享同一个 IR。

Codegen 逐步指南

1. 识别 op

在进行前几个 codegen 时，我们通常建议从更简单的 op 开始。本指南将以一个一元 op 和一个二元 op 为例，但建议您避免使用具有以下特征的 op：1. 包含自定义回退代码。例如，在 _adaptive_avg_pool3d 中，有一个条件回退。

if (!IsSupportedAdaptivePool(XlaHelpers::I64List(self.sizes()),
                             output_size_list, /*pool_dim=*/3)) {
  return at::native::call_fallback_fn<&xla_fallback, ATEN_OP(_adaptive_avg_pool3d)>::call(self, output_size);
}

  1. 导致动态形状，因为这些 op 尚在开发中，并且可能会随时间而演变。在未来的某个时候，我们可能会将这些 op 引入 codegen。

  2. 不直接调用 tensor_method。例如：

if (!self_tensor) {
  static bool sync_update =
      torch_xla::runtime::sys_util::GetEnvBool("XLA_TENSOR_UPDATE_SYNC", true);
  XLA_CHECK(dst_tensor);
  dst_tensor->UpdateFromTensor(self, /*sync=*/sync_update);
}

  1. 具有复杂的 tensor_method，理想情况下它应该是一个从 op 到 IR 的直接映射。

一个很好的“简单” op 示例将是 abs 这样的操作。

at::Tensor XLANativeFunctions::abs(const at::Tensor& self) {
  TORCH_LAZY_FN_COUNTER("xla::");
  return bridge::AtenFromXlaTensor(XLATensor::abs(bridge::GetXlaTensor(self)));
}

2. Codegen op 并检查生成的文件

在 xla/codegen/xla_native_functions.yaml 中找到该 op，将其移至 full_codegen 列，然后再次在 xla 目录下运行 python setup.py install。构建将失败（原因在本指南后面解释），但您仍然可以看到生成的文件。

如果在生成文件时遇到与 shape_inference.h 相关的错误，则可能是 PyTorch 尚未为要生成的功能提供必要的实现。您可以尝试在 shape_inference.h 中添加必要的功能以解除阻塞。

下面的代码片段使用 abs 作为示例。#### XLANativeFunctions.cpp

at::Tensor XLANativeFunctions::abs(const at::Tensor & self) {
  TORCH_LAZY_FN_COUNTER("xla::");
  auto common_device = torch_xla::bridge::GetXlaDevice(self);
  TORCH_INTERNAL_ASSERT(common_device);

  torch_xla::XLATensorPtr lazy_self = torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(self, *common_device);

  torch::lazy::NodePtr node = torch::lazy::ReuseNode<Abs>(lazy_self->GetIrValue());
  if (!node) {
    node = torch_xla::MakeNode<Abs>(lazy_self->GetIrValue());
    CacheNode(node);
  }

  auto result = torch_xla::bridge::AtenFromXlaTensor(
        torch_xla::XLATensor::Create(std::move(node), *common_device));
  return result;
};

逐行描述生成代码：- 获取并验证输入张量的设备

auto common_device = torch_xla::bridge::GetXlaDevice(self);
TORCH_INTERNAL_ASSERT(common_device);

检查是否可以重用先前创建的节点。如果不能，则创建相应的 IR 节点并缓存它。

torch::lazy::NodePtr node = torch::lazy::ReuseNode<Abs>(lazy_self->GetIrValue());
if (!node) {
  node = torch_xla::MakeNode<Abs>(lazy_self->GetIrValue());
  CacheNode(node);
}

将新创建的 IR 节点包装在 XLATensor 中。然后将 XLATensor 包装在 at::Tensor 中并作为结果返回。请注意，这部分以前是在 tensor_method.cpp 中手动完成的。

auto result = torch_xla::bridge::AtenFromXlaTensor(
      torch_xla::XLATensor::Create(std::move(node), *common_device));
return result;

LazyIr.h

class Abs : public XlaNode {
 public:
  Abs(const torch_xla::XlaValue& self)
      : XlaNode(torch::lazy::OpKind(at::aten::abs), {self},
                [&]() { return AbsOutputShape(self); },
                /* num_outputs */ 1, torch::lazy::MHash())
  {}

  std::string ToString() const override {
    std::stringstream ss;
    ss << XlaNode::ToString();
    return ss.str();
  }
  torch_xla::XlaOpVector Lower(LoweringContext* loctx) const override;
};

需要注意的几点：- Codegen 不会生成预期的 Clone 方法。即使在当前的 PyTorch/XLA 中也没有使用 Clone 方法，我们将在迁移过程中删除它们。- 对于每个 op，它将生成一个 {OP}OutputShape 方法。我们需要在单独的文件中手动声明和实现此方法。- 对于每个 op，它将生成一个 Lower 声明。我们需要在单独的文件中手动实现此降级函数。

3. 实现缺失的 IR 函数

torch_xla/csrc/ops/ops_xla_shape_fn.h

声明 {OP}OutputShape。

xla::Shape AbsOutputShape(const XlaValue& input);

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

实现 {OP}OutputShape。

xla::Shape AbsOutputShape(const XlaValue& input) { return input.xla_shape(); }

Abs 是一个过于简化的示例，在正常情况下，您需要再次调用 BuildXXXOp 函数来获取输出形状。一个稍微好一点的例子是：

xla::Shape MaximumOutputShape(const XlaValue& input, const XlaValue& other) {
  auto lower_for_shape_fn =
      [&](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
    auto promoted = XlaHelpers::Promote(operands[0], operands[1]);
    return xla::Max(promoted.first, promoted.second);
  };
  return InferOutputShape({input.xla_shape(), other.xla_shape()},
                          lower_for_shape_fn);
}

请注意，您不应从头开始。从现有 op 中查找 Xla::Shape 计算逻辑，并将其移至这两个文件。

4. 实现降级函数

torch_xla/csrc/ops/ops_lower_fn.cpp

torch_xla::XlaOpVector Abs::Lower(LoweringContext* loctx) const {
  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
  return ReturnOp(BuildAbs(xla_input), loctx);
}

请注意，此函数应直接从现有的降级函数迁移。一些最初在 torch_xla/csrc/ops/ops.cpp 中实现的 op 使用 GenericOp。您需要对其降级实现进行少量修改，以使其符合上面提供的实现。

5. 清理

从 aten_xla_type.cpp、tensor_methods.h、tensor_methods.cpp 和 ops/… 中删除现有的 op。请注意，有时您必须保留 tensor_method，因为它被用于 tensor_ops，例如。因此，在删除 op 之前，请将其与 tensor_ops.cpp 进行交叉引用。

XLATensor s1 = XLATensor::sub(XLATensor::mul(u2, v3), XLATensor::mul(u3, v2), one);

有时其他 IRNode 使用您迁移的“IRNode”。在这种情况下，您还需要更新那些 IRNode 的降级逻辑。从长远来看，我们需要从我们这边摆脱这些复合 IR，并为每个 op 提供降级函数。

torch::lazy::NodePtr exp = Pow(Abs(input), norm_exp);

至

torch::lazy::NodePtr exp =
    Pow(torch_xla::MakeNode<Abs>(input, std::vector<torch::lazy::Shape>()),
        norm_exp);

运行测试并验证结果

运行 C++ op 测试或仅包含已生成 op 的简单测试。要运行 C++ 测试：1. 通过 python setup.py install 构建 xla（注意：不要使用 BUILD_CPP_TESTS=0 标志，因为这会跳过构建 C++ 测试）2. 进入 pytorch/xla 中的 test/cpp/build 目录3. 运行命令以运行所需的 C++ 测试（例如，要运行 Abs C++ 测试）

./test_ptxla --gtest_filter=AtenXlaTensorTest.TestAbs

照常，需要验证的两件事是正确性以及 xla 计数器是否正确递增。

示例 PR

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源

为了分析流量并优化您的体验，我们在此网站上使用 Cookie。点击或浏览即表示您同意我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括可用的控件： Cookie 政策。