• 文档 >
  • 核心 ATen 算子集定义
快捷方式

核心 ATen 算子集定义

本页介绍了核心 ATen 算子集 (opset) 的描述和背景信息。本页推荐给正在为 ExecuTorch 开发新内核库或委托的开发者阅读。此外,建议先熟悉 torch.export 作为前置条件;特别是要了解 torch FX 图、算子分解和函数化的概念。

被确定为核心 ATen 算子的列表可在 PyTorch 文档网站的 IRs 页面找到。

什么是算子集?

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 库中不仅定义了大量算子,而且可能会添加新算子,或现有算子的 schema 可能会发生变化。如果没有算子分解,构建在 torch.export 生成的 IR 之上的后端将不得不处理大量的算子以及不断变化的 opset。核心 ATen opset 通过定义一个更小、更易于管理的算子集来解决这个问题,该算子集在开发时考虑了稳定性。

核心 ATen 算子集的开发

尽管 ExecuTorch 使用核心 ATen opset,但它并非 ExecuTorch 独有。核心 ATen opset 的主要设计目标之一是尽可能通用;绝大多数用例不会希望分解其中包含的算子。因此,核心 ATen opset 所隐含的分解应该对绝大多数用例有用。

另一个关键考虑因素是尽可能保持 opset 最小化,但不能以强加对性能或开发者体验产生深远负面影响的分解为代价。

核心 ATen opset 是通过调研公共 GitHub 仓库中的模型以及知名开源模型,梳理出 ATen 算子列表后进行审查而开发的。调研过程的目的是获得一个精简的 ATen 算子列表,该列表可以大致反映哪些 ATen 算子使用最广泛。这样就可以优先审查最常用的算子。

关于每个算子是应该作为核心算子还是由核心 ATen 分解表进行分解的决定是根据以下因素确定的:

  1. 检查算子的潜在分解;分解应该是使用其他 ATen 算子对该算子进行的相对直接的重新表达。

    • 分解不应该看起来像是对算子的完全实现。

    • 分解不应该随输入的运行时特性而变化。

    • 我们还考虑分解该算子是否会影响输出的精度、数值有效性或内存布局。

  2. 考虑开发者是否会出于性能或其他原因希望在图中保留该算子。

    • 例如,某个算子或许可以分解,但在大多数平台上它可以映射到单个硬件指令,在这种情况下,最好将其提升为核心算子。

未来工作

在审查并指定每个 ATen 算子为“核心”或“默认分解”之前,核心 ATen opset 不能被视为完全完成。然而,这是一项艰巨的任务,并且存在大量不常用算子的长尾部分。正因如此,才采取了调研模型以确定哪些算子最常用的方法,从而可以优先处理“更高影响”的算子。

尽管如此,仍有许多尚未评估的算子。计划是随着需求的出现继续评估更多算子;PyTorch 社区可以通过创建 GitHub Issue 或在 PyTorch 论坛上的这篇文章中 评论 来提议更多的核心算子或核心分解。

文档

查阅 PyTorch 的完整开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源