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

Codegen 迁移指南

随着 PyTorch/XLA 迁移到 LTC（Lazy Tensor Core），我们需要清理用于执行 op 下沉的现有桩代码（涵盖 6 个以上文件）。旧 op 下沉的完整过程和文件结构可以在 op 下沉指南中找到：ref:'op-lowering'。用 codegen 替换受支持的 op 不应引入任何新行为，这纯粹是为了清理。

开始之前

你应该遵循这里的说明来安装所需的依赖项并从源代码构建 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 后端使用的现有数据模型和辅助函数构建，并添加了惰性张量后端特有的新功能。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。

  • 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 这样的 op

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。构建会失败（原因将在本指南后面解释），但你仍然可以看到生成的文件。以下代码片段以 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 的 Cookies 政策适用。了解更多信息，包括可用的控制选项：Cookies 政策。