核心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 操作符的功能化 ATen 操作集。
核心 ATen 操作符集¶
导出的图可以通过应用算子分解进一步转换。此过程将用其他ATen算子的等效序列替换指定的ATen算子。例如,aten.hardsigmoid 可被替换为 aten.clamp(aten.clamp(self + 3, min=0), max=6) / 6。
如果使用默认的分解设置对 PyTorch 程序进行分解,则生成的图 IR 将包含“核心 ATen”操作集。该操作集将是功能 ATen 操作集的一个子集,因为部分操作会被分解。属于核心 ATen 操作集的操作(即核心 ATen 操作)在默认分解设置下不会被分解。通常情况下,核心 ATen 操作无法通过其他 ATen 操作的分解轻松重表达。
核心ATen操作集的主要动机在于,减少模型导出后PyTorch后端和编译器需要处理的操作符数量。ATen库中定义了大量操作符,而且未来还可能新增操作符,或现有操作符的签名发生变化。如果没有操作符分解,基于torch.export生成的中间表示(IR)构建的后端将不得不同时应对庞大的操作符集合以及不断变化的操作集。核心ATen操作集通过定义一个更小、更易于管理的操作符集合来解决这一问题,该集合的设计以稳定性为首要考虑。
核心ATen算子集的开发¶
虽然 ExecuTorch 使用了核心 ATen 算子集,但它并非专为 ExecuTorch 设计。核心 ATen 算子集的主要设计目标之一是尽可能通用;绝大多数使用场景并不希望分解其中包含的算子。因此,由核心 ATen 算子集所隐含的分解方案应对绝大多数使用场景都具备实用性。
另一个关键考量是尽可能保持操作集(opset)精简,但绝不能以牺牲性能或开发者体验为代价,强行引入会导致严重负面影响的分解操作。
核心 ATen 算子集是通过审查公开 GitHub 仓库中的模型以及知名开源模型所创建的 ATen 算子列表而开发的。该调查过程的目的是获得一个精简的 ATen 算子列表,作为哪些 ATen 算子使用频率最高的代理指标。通过这种方式,可以优先审查最常用的算子。
每个算子应作为核心算子,还是通过核心 ATen 分解表进行分解,这一决定由以下因素确定:
检查算子的潜在分解方式;该分解应是通过其他 ATen 算子对原算子进行相对直接的重新表达。
分解不应看起来像是对算子的直接实现。
分解不应基于输入的运行时特征而变化。
我们还考虑算子分解是否会影响输出的精度、数值有效性或内存布局。
考虑开发人员是否希望出于性能或其他原因在计算图中保留该算子。
例如,某个算子或许可以被分解,但在大多数平台上它可以映射为单条硬件指令,在这种情况下,将其提升为核心算子更为可取。
未来工作¶
在每一个 ATen 算子都经过审查并被指定为“核心”或“默认分解”之前,核心 ATen 算子集尚不能被视为完全完备。然而,这是一项艰巨的任务,且存在大量使用频率较低的算子长尾。因此,我们采取了一种方法:通过调研模型来确定哪些算子最为常用,从而优先处理那些“影响更大”的算子。
尽管如此,仍有许多运算符尚未经过评估。计划是根据需要继续评估其他运算符;PyTorch 社区可以通过在 GitHub 上提交问题或在 PyTorch 论坛上评论这篇文章来提议其他核心运算符或核心分解方法。