Relay中间表示介绍

最新推荐文章于 2025-01-20 01:40:40 发布

牛牛存

最新推荐文章于 2025-01-20 01:40:40 发布

阅读量2k

点赞数 1

分类专栏： tvm TVM Relay

TVM Relay 同时被 2 个专栏收录

25 篇文章

订阅专栏

tvm

14 篇文章

订阅专栏

介绍

这篇文章介绍Relay——第二代NNVM。我们期望的读者有两类背景，有编程语言开发经历和对计算图表示熟悉的深度学习框架开发人员。

我们在这里简要总结设计目标，并将在本文的后半部分涉及这些要点。

支持传统的数据流样式的编程和转换。
支持功能样式的作用域，绑定并使其成为功能齐全的可区分语言。
能够允许用户混合两种编程样式。

用Relay构建一个计算图

传统的深度学习框架使用计算图作为其中间表示。计算图（或数据流图）是代表计算的有向无环图（DAG）。尽管数据流图由于缺乏控制流而在计算能力方面受到限制，但它们的简单性使其易于实现自动微分以及针对异构执行环境进行编译（例如，在专用硬件上执行图的某些部分）。

您可以使用Relay来构建计算（数据流）图。具体来说，以上代码显示了如何构造简单的两节点图。您可以发现该示例的语法与现有的计算图IR（例如NNVMv1）没有太大区别，唯一的区别是术语：

现有框架通常使用图和子图
Relay使用函数，例如【fn(%x】)，来表示图

每个数据流节点都是Relay中的一个CallNode。Relay Python DSL允许您快速构建数据流图。我们要在上面的代码中强调的一件事——我们显式构造了一个Add节点，两个输入点都为%1。当深度学习框架评估上述程序时，它将按拓扑顺序计算节点，并且%1仅计算一次。尽管对于深度学习框架构建者来说这是很自然的事情，但它首先会让PL研究人员感到惊讶。如果我们实现一个简单的访问操作以打印出结果并将结果视为嵌套的Call表达式，则它将变为【log(%x)+log(%x)】

当DAG中存在共享节点时，这种歧义是由程序语义的不同解释引起的。在正常的函数式编程IR中，嵌套表达式被视为表达式树，而没有考虑%1实际在中重复使用的事实%2。

Relay IR注意这一区别。通常，深度学习框架用户以这种方式构建计算图，其中经常发生DAG节点重用。结果，当我们以文本格式打印Relay程序时，我们每行打印一个CallNode并为每个CallNode 分配一个临时ID【(%1, %2)】，以便可以在程序的后续部分中引用每个公共节点。

模块：支持多种函数（图）

到目前为止，我们已经介绍了如何构建作为函数的数据流图。一个人自然会问：我们可以支持多种函数并使它们彼此调用吗？Relay允许在一个模块中将多个函数组合在一起；下面的代码显示了一个函数调用另一个函数的示例。

def @muladd(%x, %y, %z) {
  %1 = mul(%x, %y)
  %2 = add(%1, %z)
  %2
}
def @myfunc(%x) {
  %1 = @muladd(%x, 1, 2)
  %2 = @muladd(%1, 2, 3)
  %2
}

该模块可以视为【Map<GlobalVar, Function>】。在这里，GlobalVar只是一个ID，用于表示模块中的函数。和在上面的例子中【@muladd】和【@myfunc】是GlobalVars。当使用CallNode调用另一个函数时，相应的GlobalVar存储在CallNode的op字段中。它包含一个间接级别——我们需要使用相应的GlobalVar从模块中查找被调用函数的主体。在这种情况下，我们还可以将对函数的引用直接存储为CallNode中的op。那么，为什么我们需要引入GlobalVar？主要原因是GlobalVar解耦了定义/声明，并启用了函数的递归和延迟声明。

def @myfunc(%x) {
  %1 = equal(%x, 1)
   if (%1) {
      %x
   } else {
     %2 = sub(%x, 1)
     %3 = @myfunc(%2)
      %4 = add(%3, %3)
      %4
  }
}

在上面的示例中，【@myfunc】递归调用自身。使用GlobalVar 【 @myfunc】表示函数可以避免数据结构中的循环依赖性。至此，我们已经介绍了Relay中的基本概念。值得注意的是，与NNVMv1相比，Relay具有以下改进：

简洁的文本格式，简化了编写过程的调试。
在联合模块中，对子图功能的一流支持可为联合优化（例如内联和调用约定规范）提供更多机会。
朴素的前端语言交互操作，例如，所有数据结构都可以在Python中访问，这允许在Python中快速优化的原型并将其与C ++代码混合。

让绑定和作用域

到目前为止，我们已经介绍了如何以深度学习框架中使用的旧方法来构建计算图。本节将讨论Relay引入的新的重要构造-let绑定。

在每种高级编程语言中都使用Let绑定。在Relay中，它是具有三个字段的数据结构【Let(var, value, body)】。在评估let表达式时，我们首先评估值部分，将其分配给var，然后在主体表达式中返回评估结果。

您可以使用一系列的let绑定来构造一个逻辑上等效于数据流程序的程序。下面的代码示例显示一个程序并排两种形式。

嵌套的let绑定称为A-normal形式，在功能编程语言中通常用作IR。现在，请仔细看一下AST的结构。尽管这两个程序在语义上是相同的（它们的文本表示形式也一样，除了A-normal形式具有let前缀），但它们的AST结构却不同。

由于程序优化采用了这些AST数据结构并对其进行了转换，因此两种不同的结构将影响我们将要编写的编译器代码。例如，如果我们要检测一个模式：【add(log(x),y)】

在数据流形式中，我们可以首先访问add节点，然后直接查看其第一个参数以查看它是否为log
在A-normal形式中，我们无法直接进行检查，因为要添加的第一个输入是【%v1】–我们将需要使映射从变量到其绑定值并查找该映射，以便知道这【%v1】是一个log。

不同的数据结构将影响您编写转换的方式，我们需要牢记这一点。因此，现在，作为深度学习框架开发人员，您可能会问：为什么我们需要let绑定？您的PL朋友总是会告诉您，let我们重要-由于PL是一个非常成熟的领域，因此背后必须有一些智慧。

为什么我们可能需要让绑定

let绑定的一种关键用法是它指定计算范围。让我们看下面的示例，该示例不使用let绑定。

当我们尝试决定应该在哪里评估node【%1】时，问题就来了。特别是，虽然文本格式似乎建议我们应该评估【%1】范围之外的节点，但AST（如图所示）却不建议这样做。实际上，数据流图永远不会定义其评估范围。这在语义上引入了一些歧义。

当我们有闭包时，这种歧义变得更加有趣。考虑下面的程序，该程序返回一个闭包。我们不知道应该在哪里计算【%1】; 它可以在封闭内部或外部。

fn (%x) {
  %1 = log(%x)
  %2 = fn(%y) {
    add(%y, %1)
  }
  %2
}

let绑定解决了这个问题，因为值的计算发生在let节点上。在两个程序中，如果我们更改【%1 = log(%x)】为【let %v1 = log(%x)】，我们都将计算位置明确指定为if范围和闭包之外。如您所见，let-binding为计算站点提供了更为精确的规范，并且在我们生成后端代码时很有用（因为此类规范在IR中）

另一方面，没有指定计算范围的数据流形式也有其自身的优势——即，我们无需担心在生成代码时将let放在哪里。数据流形式还为以后的遍历提供了更大的自由度，以决定将评估点放在何处。因此，如果发现方便，在优化的初始阶段使用程序的数据流形式可能不是一个坏主意。今天在Relay中进行了许多优化，以优化数据流程序。

但是，当我们将IR降低到实际的运行时程序时，我们需要精确计算范围。特别是，当我们使用子函数和闭包时，我们要明确指定计算范围应在哪里发生。在后期执行特定的优化中，可以使用let绑定来解决此问题。

IR转换的意义

希望到目前为止，您已经熟悉两种表示形式。大多数函数式编程语言均以A-normal进行分析，分析人员无需注意表达式是DAG。

Relay选择同时支持数据流形式和let绑定。我们认为让框架开发人员选择他们熟悉的表示形式很重要。但是，这确实对我们编写通行证有一些影响：

如果您来自数据流背景并且想要处理let，请保留var到表达式的映射，以便在遇到var时可以执行查找。这可能意味着最小的更改，因为无论如何我们已经需要从表达式到转换后的表达式的映射。请注意，这将有效删除程序中的所有let。
如果您来自PL背景并且喜欢A-normal格式，我们将为A-normal格式传递提供数据流。
对于PL人员来说，当您实现某种东西（例如，从数据流到ANF的转换）时，请注意表达式可以是DAG，这通常意味着我们应该访问带有【Map<Expr, Result>】表达式，并且只计算一次转换后的结果，因此结果表达式保持通用结构。Map<Expr, Result>

还有一些其他高级概念，例如符号形状推断，多态函数等，本材料未介绍这些概念。非常欢迎您查看其他材料。