[翻译]Relay IR简介— tvm 0.6.0文档

Relay IR简介

本文介绍了Relay IR——第二代NNVM。我们希望读者具有两种背景:具有编程语言背景的读者和熟悉计算图表示的深度学习框架开发人员。简要总结设计目标如下,在本文的后半部分将会涉及这些要点。

  • 用Relay建立计算图

    传统的深度学习框架使用计算图作为其中间表示。计算图(或数据流图)是代表计算的有向无环图(DAG)。尽管数据流图由于缺乏控制流而在计算能力方面受到限制,但它们的简单性使其易于实现自动区分并针对异构执行环境进行编译(例如,在专用硬件上执行图的某些部分)。[翻译]Relay IR简介— tvm 0.6.0文档 - 图1你可以使用Relay来构建计算(数据流)图。具体来说,以上代码显示了如何构造简单的两节点图。你会发现该示例的语法与现有的计算图形IR(例如NNVMv1)没有太大区别,唯一的区别是术语:
  • 现有框架通常使用图和子图
  • Relay 使用例如fn(%x)的函数来表示图 每个数据流node都是Relay中的一个CallNode。通过Relay Python DSL,你可以快速构建数据流图。在上面的代码中值得强调的一点是我们显式构造了一个Add节点,两个输入点都为%1。当深度学习框架评估上述程序时,它将按拓扑顺序计算节点,并且%1仅计算一次。虽然这对于深度学习框架构建者来说是很自然的事情,但它首先会让PL研究人员感到惊讶。如果我们实现一个简单的访问者以打印出结果并将结果视为嵌套的Call表达式,则它将变为。log(%x) + log(%x)当DAG中存在共享节点时,这种歧义是由程序语义的不同解释引起的。在正常的函数式编程IR中,嵌套表达式被视为表达式树,而没有考虑%1实际在%2中重复使用的事实。Relay IR注意这一区别。通常,深度学习框架用户以这种方式构建计算图,其中经常发生DAG节点重用。结果,当我们以文本格式(Text Form)打印Relay程序时,我们每行打印一个CallNode并为每个CallNode分配一个临时ID,以便可以在程序的后续部分中引用每个公共节点。(%1,%2)## 模块:支持多种函数(图形) 到目前为止,我们已经介绍了如何构建数据流图作为函数。一个自然而然的问题是:我们可以支持多种函数并使它们相互调用吗?Relay 允许将多个函数组合在一个模块中;下面的代码显示了一个函数调用另一个函数的示例。该模块可以视为Map,在这里,GlobalVar只是一个ID,用于表示模块中的函数。@muladd和@myfunc是上面的例子中的GlobalVars。当使用CallNode调用另一个函数时,相应的GlobalVar存储在CallNode的op字段中。它包含一个间接的层次——我们需要使用相应的GlobalVar从模块中查找被调用函数的主体。在这种情况下,我们还可以将对函数的引用直接存储为CallNode中的op。那么,为什么我们需要引入GlobalVar?主要原因是GlobalVar解耦了定义/声明,并启用了函数的递归和延迟了函数的声明。在上面的示例中,@myfunc递归调用自身。使用GlobalVar@myfunc表示函数可以避免数据结构中的循环依赖性。至此,我们已经介绍了Relay中的基本概念。值得注意的是,与NNVMv1相比,Relay 具有以下改进:
  • 简洁的文本格式,可简化书写Pass的调试。
  • 在联合模块中,对子图函数的一流支持可为联合优化(例如内联和调用约定规范)提供更多机会。
  • 朴素的前端语言互操作,例如,所有数据结构都可以在Python中访问,这允许在Python中快速优化的原型并将它们与C++代码混合。

    到目前为止,我们已经介绍了如何以深度学习框架中使用的旧方法来构建计算图。本节将讨论由Relay引入的新的重要构造- Let绑定。在每种高级编程语言中都使用 Let绑定。在Relay中,它是具有三个字段的数据结构Let(var,value,body)。在评估Let表达式时,我们首先评估value部分,将其分配给var,然后在body表达式中返回评估结果。你可以使用一系列的 Let绑定来构造一个逻辑上等效于数据流程序的程序。下面的代码示例显示一个程序的两种形式。[翻译]Relay IR简介— tvm 0.6.0文档 - 图2嵌套的Let绑定称为A-normal形式,在功能编程语言中通常用作IR。现在,请仔细看一下AST的结构。尽管这两个程序在语义上是相同的(它们的文本表示形式也一样,除了A-normal形式具有let前缀),但它们的AST结构却不同。由于程序优化采用了这些AST数据结构并对其进行了转换,因此两种不同的结构将影响我们将要编写的编译器代码。例如,如果我们要检测一个模式:add(log(x), y)
  • 在数据流形式中,我们可以首先访问添加的节点,然后直接查看其第一个参数以查看它是否为日志
  • 在A-normal形式中,我们无法直接检查,因为要添加的第一个输入是%v1–我们将创建从变量到其绑定值的映射,并查找该映射,以便知道这%v1是一个日志。 我们要牢记这一点:不同的数据结构将影响你编写转换的方式。因此,现在,作为深度学习框架开发人员,你可能会问:为什么我们需要 Let绑定?你的PL朋友总是会告诉你,let非常重要——由于PL是一个非常成熟的领域,因此背后必须有一些智慧。## 为什么我们需要let绑定 Let绑定的一种主要用法是它指定计算范围。让我们看下面的示例,该示例不使用 Let绑定。[翻译]Relay IR简介— tvm 0.6.0文档 - 图3当我们尝试决定应该在哪里评估节点 %1时,问题就来了。特别是,虽然文本格式似乎建议我们应该在if范围之外评估节点%1,但是AST(如图所示)却不建议这样做。实际上,数据流图永远不会定义其评估范围。这在语义上引入了一些歧义。当我们有闭包时,这种歧义变得更加有趣。考虑下面的程序,该程序返回一个闭包。我们不知道应该在哪里计算%1,它可以在闭包的内部或外部。fn(%x){ %1=log(%x) %2=fn(%y){ add(%y,%1) %2Let绑定解决了这个问题,因为value的计算发生在let节点上。在上面两个程序中,如果将%1=log(%x)改为let%v1=log(%x),明确将计算位置指定为if范围和闭包之外。如你所见,let绑定为计算站点提供了更为精确的规范,并且在我们生成后端代码时很有用(因为IR中有此类的规范)。另一方面,没有指定计算范围的数据流形式也有其自身的优势——我们无需担心在生成代码时将let放在哪里。数据流形式还为以后的pass提供了更大的自由,以决定将评估点放在何处。因此,如果方便的话,在优化的初始阶段使用程序的数据流形式可能是一个好主意。现在在Relay中专门做了许多优化,用来优化数据流程序。但是,当我们将IR运用到实际的运行时程序时,我们需要精确计算范围。特别是,当我们使用子函数和闭包时,我们要明确指定计算范围应在哪里发生。在后期执行特定的优化中,可以使用 Let绑定来解决此问题。## IR转换的意义 希望到目前为止,你已经熟悉两种表示形式。大多数函数式编程语言都以A-normal形式进行分析,分析人员无需注意表达式是DAG。Relay 选择支持数据流形式和let绑定。我们认为让框架开发人员选择他们熟悉的表示形式很重要。但是,这确实对我们编写pass有一些影响:
  • 如果你有数据流背景并且想要处理let,请保留var到表达式的映射,以便遇到var时可以执行查找。这可能意味着最小的更改,因为无论如何我们需要从表达式到转换后的表达式的映射。请注意,这将有效删除程序中的所有let。
  • 如果你来自PL背景并且喜欢A-normal形式,我们将为A-normal形式的pass提供一个数据流。
  • 对于PL人员来说,当你实现某种东西(例如,从数据流到ANF的转换)时,请注意表达式可以是DAG,这通常意味着我们应该访问带有一个Map的表达式,并且只计算一次转换后的结果,因此结果表达式保持通用结构。 还有其他高级概念,例如符号形状推断,多态函数等,本文未涵盖;欢迎你查看其它材料。