概述

基本介绍

本章介绍自定义PyTorch算子如何与TorchAir图模式配合工作（也叫自定义算子入图）。阅读本节，您将了解自定义PyTorch算子与TorchAir配合的完整流程，以及流程涉及的交付件及其作用。

TorchAir是昇腾为PyTorch图模式（torch.compile）实现的NPU后端，提供两种工作模式（通过config.mode切换）：

• reduce-overhead模式，又称为aclgraph模式，提供模型下沉调度、多Stream并行、图间内存复用等能力。aclgraph模式与Torch原生图模式所需的交付件完全一致。
• max-autotune模式，又称为Ascend IR模式，在aclgraph能力基础上额外提供SuperKernel等JIT编译能力，进一步提升执行性能。使用Ascend IR需要实现Ascend Converter交付件，完成将PyTorch算子转换为Ascend IR。

实现一个能与TorchAir图模式配合工作的PyTorch算子最多包含如下操作步骤，以实现In-place类算子为例。

• 步骤1：PyTorch Eager模式调用交付件，确定目标PyTorch算子原型，并完成Schema定义。
• 步骤2：PyTorch Eager模式调用交付件，基于Ascend C完成目标PyTorch算子的NPU实现。
• 步骤3：PyTorch Eager模式调用交付件，基于OpPlugin完成Ascend C算子Eager模式适配，实现Eager模式下调用运行在NPU上的PyTorch算子。
• 步骤4：PyTorch原生入图操作，完成Meta符号化推导。完成后可使用TorchAir reduce-overhead模式（aclgraph模式）以及aot_eager等原生图模式后端，并获取基于aclgraph的下沉调度收益。
• 步骤5（可选）：PyTorch原生入图操作，仅当算子为In-place类算子才需要额外完成本步骤。
• 步骤6（可选）：Ascend IR入图操作，如果您希望使用SuperKernel、分核执行等TorchAir max-autotune模式提供的高阶能力，才需要额外完成本步骤。

• 本章中提到的Eager模式、算子Schema、Ascend C、OpPlugin、In-place算子等术语请先参考常用概念进行了解，避免概念混淆。
• 对于图中的“Meta推导函数”，PyTorch原生要求所有能与torch.compile配合工作的算子需要实现Meta推导函数，又称为“符号化推导”。Meta函数表示了PyTorch算子输出与输入shape、dtype以及内存的关系，它是PyTorch入图的前提条件，其详细介绍请参考PyTorch官网符号化手册。
• 对于图中的“函数化”，请您先了解In-place与非In-place操作：
  • 计算时会修改输入的算子称为In-place算子（或原地算子、Ref类算子等），例如torch.ops.aten.add_。
  • 与torch.ops.aten.add_对应的非In-place算子为torch.ops.aten.add，其结果写入输出而非直接修改输入。

  “函数化转换”可以简单理解为将In-place算子替换为非In-place算子的过程，例如将torch.ops.aten.add_替换为torch.ops.aten.add。PyTorch图模式基于函数化后的FX图工作，因此In-place类算子与PyTorch图模式配合工作时，需要实现函数化转换，实现将图上的In-place算子替换为非In-place算子。

• 注意In-place与非In-place算子的交付件略有不同：非In-place算子本身是函数化的，无需实现函数化转换。In-place类算子需要实现函数化，TorchAir将在未来版本支持PyTorch社区的自动化函数化能力，目前仍需要您手动实现函数化。

场景说明

对于不同的入图场景，所需实现的步骤不同。

• 场景1：当您未完成自定义算子开发，希望开发Eager模式和图模式下都能工作的PyTorch算子。

  您需要完成图中步骤1~步骤6，“可选”步骤请按需实现，详细的操作参考算子开发和适配TorchAir图模式。

• 场景2：当您已完成自定义算子开发且能在Eager模式下正常工作，希望完成后续的PyTorch图模式适配。

  您需要完成图中步骤4~步骤6，“可选”步骤请按需实现，详细的操作参考算子插件化适配TorchAir图模式。

  该入图方式称为“插件化适配”，过程中无需编译、安装torch_npu，并且采用全Python实现，您可以在任意py文件中实现上述交付件，并在模型执行前加载即可。方便进行算子调试或者将自定义算子模块作为插件使用。

样例说明

针对不同的入图场景、不同的算子类型（In-place/非In-place算子），本章提供了对应的开发示例（罗列关键步骤，请根据实际业务场景调整示例代码）：

• 场景1：当您未完成自定义算子开发，希望开发Eager模式和图模式下都能工作的PyTorch算子。
  • 非In-place算子开发和入图样例
  • In-place算子开发和入图样例
• 场景2：当您已完成自定义算子开发且能在Eager模式下正常工作，希望完成后续的PyTorch图模式适配。
  • 非In-place算子插件化入图样例
  • In-place算子插件化入图样例