核心 ATen 运算符集的定义¶
本页提供了 Core ATen Operator Set (opset) 的描述和背景。建议为 ExecuTorch 开发新内核库或委托的人阅读本页。还建议将 熟悉 作为先决条件;特别是 torch FX 图、运算符分解和函数化的概念。
已在 PyTorch 文档网站的 IR 页面上找到已被确定为核心 ATen 运算符的运算符列表。
什么是运算符集?¶
torch.export
在给定的 PyTorch 程序上执行完整的图形捕获,生成描述程序执行的计算的图形 IR。运算符(即对 Tensor 执行的操作)是图中的基本计算单位,通常对应于图 IR 中的唯一节点。运算符的主要来源是 ATen 库;在 ATen 算子之外,开发者也可以定义自己的算子(即自定义算子)。
“ATen 运算符集”或“ATen opset”是一组 ATen 运算符,一旦 PyTorch 程序被捕获到图形 IR 中,就可以用来表示它。
功能性 ATen 运算符集¶
的程序捕获机制 会生成一个函数化的图,它只允许函数式运算符(即不改变或别名输入的运算符)。因此,将生成一个图形,该图形将包含函数式 ATen 运算集,而函数式 ATen 运算集仅包含函数式 ATen 运算符。torch.export
torch.export
核心 ATen 运算符集¶
可以通过应用运算符分解来进一步转换导出的图形。此过程会将指定的 ATen 运算符替换为其他 ATen 运算符的等效序列。例如,可以替换为 .aten.hardsigmoid
aten.clamp(aten.clamp(self + 3, min=0), max=6) / 6
如果使用默认分解设置对 PyTorch 程序进行分解,则生成的图形 IR 将包含“核心 ATen”操作集。此 opset 将是函数式 ATen opset 的子集,因为某些运算符将被分解。作为核心 ATen opset 一部分的 ATen 运算符(即核心 ATen 运算符)在默认分解设置下不会分解。通常,核心 ATen 算子不能轻易地通过分解被其他 ATen 算子重新表示。
核心 ATen opset 背后的主要动机是减少导出模型后 PyTorch 后端和编译器需要处理的运算符数量。不仅 ATen 库中定义了大量运算符,而且可能会添加新的运算符,或者现有运算符的架构可能会更改。如果不分解算子,构建在 IR 之上的后端将不得不处理一个大型算子表面,以及一个不断变化的 opset。核心 ATen opset 通过定义一组更小、更易于管理的运算符来解决此问题,该运算符在开发时考虑了稳定性。torch.export
核心 ATen 算子集的开发¶
尽管 ExecuTorch 使用核心 ATen opset,但它并不特定于 ExecuTorch。核心 ATen opset 的主要设计目标之一是它应该尽可能通用;绝大多数用例都不想分解其中包含的运算符。推而广之,核心 ATen opset 隐含的分解应该对绝大多数用例有用。
另一个关键考虑因素是使 opset 尽可能少,但不能以施加对性能或开发人员体验产生深远负面影响的分解为代价。
核心 ATen opset 是通过查看通过调查公共 GitHub 存储库中的模型以及众所周知的开源模型而创建的 ATen 运算符列表来开发的。调查过程的目的是获得 ATen 运算符的简化列表,该列表是使用最多的 ATen 运算符的代理。这样,可以首先查看最常用的运算符。
决定每个算子是核心算子还是由核心 ATen 分解表分解,由以下因素决定:
检查运算符的潜在分解;分解应该是使用其他 ATen 算子对算子进行相对简单的重新表达。
分解不应看起来像 Operator 的直接实现。
分解不应根据输入的运行时特征而变化。
我们还考虑了分解运算符是否会影响输出的精度、数值有效性或内存布局。
考虑开发人员是否出于性能或其他原因希望在图形中保留运算符。
例如,也许一个 operator 可以被分解,但它可以在大多数平台上映射到单个硬件指令,在这种情况下,最好将其提升为核心 operator。
未来的工作¶
在审查每个 ATen 运算符并指定“核心”或“默认分解”之前,核心 ATen 运算集不能被视为完全完成。然而,这是一项艰巨的任务,并且有一长串不经常使用的运算符。这就是为什么采取了一种方法,对模型进行调查,以确定哪些操作是最常用的,从而允许优先考虑“影响更大”的操作员。
尽管如此,仍有许多运维尚未进行评估。计划是根据需要继续评估其他运维;PyTorch 社区可以通过打开 GitHub 问题或在 PyTorch 论坛上评论此帖子来提出额外的核心运算符或其他核心分解。