深度学习编译器入坑指北(1)高级中间表示 The high-level IR

为了克服传统编译器中所采用的 IR 的局限性,避免对深度学习模型中使用的复杂计算表达的限制,现有的深度学习编译器使用经过特殊设计的高级 IR(即Graph IR,图IR,图中间表达)来进行有效的代码优化。

为了更好的理解深度学习编译器采用的图IR,后面的部分将从表示(Representation)和实现(Implementation)两部分介绍图IR 。

表示:Representation of Graph IR.

基于有向无环图的IR / DAG-based IR

基于DAG的IR是深度学习编译器最常用的计算图构造方式,在DAG中,node表示原子深度学习算子(atomic DL operators ,edge 表示张量。利用DAG,深度学习编译器可以分析不同算子之间的联系和依赖。基于DAG的IR存在计算范围定义缺失(the missing definition of computation scope)导致语义模糊(semantic ambiguity)等不足。

基于Let绑定的IR / Let-binding-based IR

let 绑定通过为某些使用范围受限的函数提供 let 表达式来解决语义模糊的问题。通过使用let 关键字,可以生产一个let node,然后表达式中指向算子和变量,而不是像DAG那样仅仅构建变量之间的计算关系。在基于DAG的编译器中,如果想要获得某个表达式的返回值,首先需要访问对应的节点,并搜索与之相关的节点(即:递归下降技术,recursive descent technique)。与之不同是,在基于let绑定的编译器中,计算出 let 表达式中变量的所有结果并构建一个变量映射(variable map),当需要某个结果时,编译器会从变量映射中查找对应的结果,以此决定表达式的结果。

张量计算的表示方式 / Representing Tensor Computation

张量计算的表示方式主要有一下几种:

1) 基于函数 Function-based

基于函数的张量计算表示,通过提供了一系列封装的算子,在算子内部实现张量计算。

2) Lambda 表达式 Lambda expression

使用 lambda 表达式可以快速定义定义计算过程,而不需要实现新的函数。

实现:Implementation of Graph IR.

数据表示 Data representation

在深度学习编译器中,通常以张量(tensor)的形式保存数据(输入,权重,中间数据等)。在深度学习编译器中可以用内存指针表示张量数据,或以占位符等更灵活的方式。一个占位符包含张量的每个维度。或者,张量的维度大小可以标记为未知。对于IR优化,深度学习编译器则需要数据布局信息。除此之外,迭代器的边界应该根据占位符来推断。

1) 占位符 Placeholder

占位符仅显式声明了形状大小,相当于定义了一个位置而不填充数据,只在计算时填充数值。这允许程序员仅仅描述计算图的操作,而无需关心具体的数据内容,这将深度学习编译器中的计算定义(the computation definition)与具体执行(the exact execution)分离。除此之外,使用占位符方便程序员更改输入、输出以及中间数据的形状,而不需要修改计算定义。

2) 动态形状表示 Dynamic shape representation

为了支持动态模型,因此需要张量的动态形状表示。声明占位符时通常用来支持动态维度大小。例如,TVM使用Any表示未知的维度信息(Tensor<(Any,3),fp32>);XLA使用None表示同样的信息(tf.placeholder("float",[None,3]))。nGraph使用PartialShape 类实现。然而,为了完全支持动态模型,应该放宽边界推断(bound inference)和维度检查。 此外,应该额外建立机制来保证内存的有效性。

3) 数据布局 Data layout

数据布局描述了张量在内存中的组织方式,通常是从逻辑索引到内存索引的映射。其通常包括维度顺序(the sequence of dimensions,即:NCHW、 NHWC等)、tiling、padding、striding 等。将数据布局信息与算子绑定而不是与张量绑定使得某些算子可以直观实现(intuitive implementation),同时可以减少编译开销。

4) 边界推断 Bound inference

在深度学习编译器进行模型编译时,边界推断被用于确定迭代器的边界。尽管 DL 编译器中的张量表示可以方便地描述输入和输出,但它在推断迭代器边界方面却面临着特殊的挑战。边界推断通常根据计算图和已知的占位符,以递归或迭代的方式执行。例如,在 TVM 中,迭代器形成一个有向无环超图(directed acyclic hyper-graph),其中图的每个节点表示一个迭代器,每个超边(hyper-edge)表示两个或多个迭代器之间的关系(例如,拆分split 、融合fuse 或变基rebase)。 一旦根据占位符的形状确定了根迭代器(root iterator)的边界,就可以根据关系递归地推断出其他迭代器。

算子支持 Operators supported

算子对应计算图中的节点,深度学习编译器支持的算子表示对应的深度学习工作。算子通常包括以下几类:

  • 代数算子(e.g., +, ×, exp and topK)
  • 神经网络算子(e.g., convolution and pooling)
  • 张量算子(e.g., reshape, resize and copy)
  • 广播(broadcast)和缩减(reduction)算子(e.g., min and argmin)
  • 控制流(control flow)算子(e.g., conditional and loop)

1) 广播算子 Broadcast

如果没有 broadcast 算子,输入张量的形状会要求更严格。例如 add 算子,要求两个输入张量的形状必须一致。使用 broadcast 算子则可以放松这个限制,通过复制输入直到形状相符,实现矩阵或向量的逐元素相加。

2) 控制流算子 Control flow

当表示复杂且灵活的模型时需要控制流。如RNN和增强学习这些依赖循环关系和具有数据依赖条件执行的模型,就需要控制流。如果图IR不支持控制流,就需要依赖主机语言(如 if 或 while)或静态展开(static unrolling)实现数据流的控制,但这样会严重损害计算效率。

最后编辑于 :
©著作权归作者所有,转载或内容合作请联系作者
【社区内容提示】社区部分内容疑似由AI辅助生成,浏览时请结合常识与多方信息审慎甄别。
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

相关阅读更多精彩内容

友情链接更多精彩内容