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

核心 ATen 算子集定义

此页面提供核心 ATen 算子集(opset)的描述和背景信息。建议正在为 ExecuTorch 开发新的内核库或委托的人员阅读此页面。还建议熟悉 torch.export 作为先决条件;特别是 torch FX 图、算子分解和函数化的概念。

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

什么是算子集?

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

“ATen 算子集”或 “ATen opset” 是 ATen 算子的集合,这些算子可用于表示 PyTorch 程序,一旦程序被捕获到图 IR 中。

函数式 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 的开发是通过审查从公共 GitHub 存储库中的模型以及著名的开源模型调查创建的 ATen 算子列表。调查过程的目的是获得一个缩减的 ATen 算子列表,该列表是使用最多的 ATen 算子的代理。这样就可以优先审查最常用的算子。

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

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

    • 分解不应看起来像对算子的彻底实现。

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

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

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

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

未来工作

在审查并指定每个 ATen 算子为 “核心” 或 “默认分解” 之前,核心 ATen opset 不能被认为是完全完整的。然而,这是一项艰巨的任务,并且存在大量不常用的算子。这就是为什么采用了一种方法,即调查模型以确定哪些 ops 是最常用的,从而可以优先处理 “影响更大” 的算子。

尽管如此,仍然有许多算子尚未评估。计划是在需要时继续评估其他算子;PyTorch 社区可以通过打开 GitHub issue 或 在 PyTorch 论坛上对此帖子发表评论 来提出额外的核心算子或额外的核心分解。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源