核心 ATen 算子集的定义¶

此页面提供核心 ATen 算子集 (opset) 的描述和背景信息。建议开发 ExecuTorch 的新内核库或委托的人员阅读此页面。此外，建议读者事先熟悉torch.export，尤其是 torch FX 图、算子分解和函数化等概念。

已识别为核心 ATen 算子的算子列表可以在PyTorch 文档网站的 IR 页面上找到。

什么是算子集？¶

torch.export 对给定的 PyTorch 程序执行完整的图捕获，生成一个描述程序执行的计算的图 IR。算子（即对张量执行的操作）是图中计算的基本单元，通常对应于图 IR 中的唯一节点。算子的主要来源是ATen 库；除了 ATen 算子之外，开发人员还可以定义自己的算子（即自定义算子）。

“ATen 算子集”或“ATen opset”是指一旦 PyTorch 程序捕获到图 IR 中，就可以用来表示该程序的 ATen 算子集。

函数式 ATen 算子集¶

torch.export 的程序捕获机制会生成一个函数化的图，该图仅允许函数式算子（即不修改或混淆输入的算子）。因此，torch.export 生成的图将包含函数式 ATen opset，其中仅包含函数式 ATen 算子。

核心 ATen 算子集¶

导出的图可以通过应用算子分解进行进一步转换。此过程将用其他 ATen 算子的等效序列替换指定的 ATen 算子。例如，aten.hardsigmoid 可以替换为 aten.clamp(aten.clamp(self + 3, min=0), max=6) / 6。

如果使用默认分解设置对 PyTorch 程序进行分解，则生成的图 IR 将包含“核心 ATen”opset。此 opset 将是函数式 ATen opset 的一个子集，因为某些算子将被分解。作为核心 ATen opset 一部分的 ATen 算子（即核心 ATen 算子）在默认分解设置下不会被分解。通常，核心 ATen 算子不能通过分解轻松地用其他 ATen 算子重新表示。

核心 ATen opset 背后的关键动机是减少导出模型后 PyTorch 后端和编译器需要处理的算子数量。ATen 库中定义的算子数量众多，而且可能会添加新的算子，或者现有算子的模式可能会发生变化。如果没有算子分解，构建在 torch.export 生成的 IR 之上的后端将不得不处理大量的算子表面，以及不断变化的 opset。核心 ATen opset 通过定义一个更小、更易于管理的算子集来解决此问题，该算子集在设计时就考虑了稳定性。

核心 ATen 算子集的开发¶

尽管 ExecuTorch 使用核心 ATen opset，但它并非特定于 ExecuTorch。核心 ATen opset 的主要设计目标之一是尽可能通用；绝大多数用例都不希望分解其中包含的算子。相应地，核心 ATen opset 所隐含的分解应该对绝大多数用例都有用。

另一个关键考虑因素是尽可能保持 opset 的最小化，但不能以牺牲会对性能或开发人员体验造成重大负面影响的分解为代价。

核心 ATen opset 是通过审查 ATen 算子列表开发的，该列表是通过调查公共 GitHub 存储库中的模型以及知名开源模型创建的。调查过程的目的是获得一个减少的 ATen 算子列表，该列表可以作为哪些 ATen 算子使用最多的代理。这样，就可以首先审查最常用的算子。

每个算子是否应成为核心算子或由核心 ATen 分解表分解的决定取决于

检查算子的潜在分解；分解应该使用其他 ATen 算子对算子进行相对简单的重新表达。
- 分解不应该看起来像是对算子的直接实现。
- 分解不应根据输入的运行时特性而变化。
- 我们还会考虑分解算子是否会影响输出的精度、数值有效性或内存布局。
考虑开发人员是否出于性能或其他原因希望在图中保留该算子。
- 例如，也许一个算子可以被分解，但它可以在大多数平台上映射到单个硬件指令，在这种情况下，最好将其提升为核心算子。

未来工作¶

在每个 ATen 算子都经过审查并被指定为“核心”或“默认分解”之前，核心 ATen 操作集都不能被认为是完全完整的。然而，这是一项艰巨的任务，并且存在大量不常使用的算子。这就是为什么采取了一种方法，即对模型进行调查以确定哪些操作是最常用的，从而可以优先考虑“影响更大”的算子。

尽管如此，仍然有许多算子尚未评估。计划是在需要时继续评估其他算子；PyTorch 社区可以通过创建一个 GitHub 问题或在 PyTorch 论坛上对此帖子发表评论来提出其他核心算子或其他核心分解。