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

Codegen 迁移指南

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

开始之前

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

export PJRT_DEVICE=CPU

还建议您在处理 codegen 之前熟悉我们的op lowering 过程。

PyTorch/XLA 使用 https://github.com/pytorch/xla/issues/3560 来跟踪 codegen 迁移的状态。在处理 codegen 时，请将您的 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。只有非结构化的 ops 才需要手动形状推断函数

  • 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。大多数 ops 都属于受支持类别，本文档的目标是将大多数 ops 移至 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 ops 的 IR。

PyTorch/XLA 旧 Op Lowering 文件

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

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

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

    • 定义 XLATensor 类和 XLATensor 方法声明。这些声明通常是我们​​在 XLANativeFunctions.h 中声明的 at::Tensor 节点的​​一对一映射。XLATensor 方法将为 full_codegen ops 删除

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

    • 实现 tensor.h 中定义的 tensor 方法。此文件将为 full_codegen ops 删除

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

    • 为“大多数” ops 定义 IR 类。多个 ops 可能共享相同的 IR。

Codegen 逐步指南

1. 识别 op

当您处理最初的几个 codegen 时，我们通常建议您从更简单的 ops 开始。本指南将介绍一个一元 op 和一个二元 op 作为示例，但建议您避免具有以下特征的 ops：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. 导致动态形状，因为这些 ops 正在进行中，并且可能会随着时间的推移而发展。在未来的某个时候，我们可能会将 ops 引入 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。构建将失败（原因在本指南后面解释），但您仍然可以看到生成的文件。以下代码片段使用 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 声明。我们需要在一个单独的文件中手动实现此 lowering 函数。

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. 实现 lowering 函数

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

请注意，此函数应直接从现有的 lowering 移动。最初在 torch_xla/csrc/ops/ops.cpp 中实现的某些 Ops 使用 GenericOp。您需要稍微修改其 lowering 实现以适应上面提供的实现。

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 lowering 逻辑。从长远来看，我们需要从我们的端点删除这些复合 IR，并为每个 op 提供一个 lowering 函数。

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 测试或仅涉及生成的 ops 的简单测试。要运行 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 政策。