核心 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 分解表进行分解的决定是根据以下因素确定的:
检查算子的潜在分解;分解应该是使用其他 ATen 算子对该算子进行的相对直接的重新表达。
分解不应该看起来像是对算子的完全实现。
分解不应该随输入的运行时特性而变化。
我们还考虑分解该算子是否会影响输出的精度、数值有效性或内存布局。
考虑开发者是否会出于性能或其他原因希望在图中保留该算子。
例如,某个算子或许可以分解,但在大多数平台上它可以映射到单个硬件指令,在这种情况下,最好将其提升为核心算子。