核心 ATen 算子集的定义¶
此页面提供了核心 ATen 算子集 (opset) 的描述和背景信息。对于那些为 ExecuTorch 开发新内核库或委托的人来说,建议阅读此页面。还建议熟悉 torch.export
作为先决条件;特别是 torch FX 图、算子分解和函数化的概念。
可在 PyTorch 文档网站的 IR 页面 上找到已识别为核心 ATen 算子的算子列表。
什么是算子集?¶
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 是通过查看在公共 GitHub 存储库中的模型以及众所周知的开源模型中创建的 ATen 运算符列表来开发的。调查过程的目的是获得 ATen 运算符的简化列表,该列表是使用最多的 ATen 运算符的代理。这样,可以首先查看最常用的运算符。
每个运算符是核心运算符还是由核心 ATen 分解表分解的决定取决于
检查运算符的潜在分解;分解应该是使用其他 ATen 运算符对运算符进行相对简单的重新表示。
分解不应看起来像运算符的直接实现。
分解不应根据输入的运行时特性而变化。
我们还会考虑分解算子是否会影响输出的精度、数值有效性或内存布局。
思考开发者是否出于性能或其他原因希望在图中保留算子。
例如,一个算子也许可以分解,但它可以在大多数平台上映射到一个硬件指令,在这种情况下,最好将其提升为核心算子。
后续工作¶
在对每个 ATen 算子进行审查并将其指定为“核心”或“默认分解”之前,核心 ATen opset 无法被认为是完全完成的。然而,这是一项艰巨的任务,并且还有很多不常用的算子。这就是采用一种方法来调查模型以确定最常用的算子,从而允许对“影响较大”的算子进行优先排序的原因。
尽管如此,仍然有很多算子尚未经过评估。计划是根据需要继续评估其他算子;PyTorch 社区可以通过打开 GitHub 问题或通过在 PyTorch 论坛上评论此帖子来提议其他核心算子或其他核心分解。