Relay IR表示
Relay IR简介
本节介绍了 Relay IR——第二代 NNVM。期待两种背景的读者——具有编程语言背景的读者和熟悉计算图表示的,深度学习框架开发人员。
简要总结了设计目标,讨论要点。
• 支持传统的数据流式编程和转换。
• 支持函数式作用域,let-binding,成为功能齐全的可微分语言。
• 能够允许用户混合两种编程风格。
使用 Relay 构建计算图
传统的深度学习框架,使用计算图作为中间表示。计算图(或数据流图)是表示计算的有向无环图 (DAG)。由于缺乏控制流,数据流图表达的计算受限制,容易实现自动微分,异构执行环境编译(例如,专用硬件上执行图)。
可以使用 Relay 构建计算(数据流)图。如何构造一个简单的二节点图。示例语法与 NNVMv1 等现有计算图 IR,唯一区别:
• 现有框架通常使用图和子图
• Relay 使用函数 eg – 表示图形fn (%x)
每个数据流节点,都是 Relay 中的一个 CallNode。Relay Python DSL快速构建数据流图。构造Add节点,两个输入点都指向%1。深度学习框架评估时,按照拓扑顺序计算节点, %1只会计算一次。实现访问者,打印结果,嵌套的 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,有以下改进:
• 简洁的文本格式,可简化写入passes的调试。
• 在联合模块中,对子图函数的一流支持,使得联合优化会成为可能,如内联和调用约定规范。
• 简单的前端语言互操作,所有数据结构,可以在 Python 中访问,允许在 Python 中,快速构建优化原型,与 C++ 代码混合。
让绑定和作用域
深度学习框架中,使用的旧方法构建计算图。将讨论 Relay 引入新结构——let 绑定。
每一种高级编程语言,使用 let 绑定。在 Relay 中,Let(var, value, body)表示三个字段的数据结构。评估let表达式,先评估 value 部分,分配给 var,在 body 表达式中,返回评估结果。
可以使用一系列 let 绑定,构造逻辑上等效数据流程序。
嵌套的 let 绑定称为 A 范式,通常用作函数式编程语言中的 IR,仔细看看 AST 结构。虽然这两个程序在语义上相同(除了 A 范式,有 let 前缀),但AST 结构不同。
程序优化采用这些 AST 数据结构,进行转换,两种不同的结构,将影响要编写的编译器代码。例如,如果检测一个模式:add(log(x),y)。
• 在data-flow形式中,先访问add节点,直接查看第一个参数,是不是日志
• 在 A 范式中,不能直接检查,要添加的第一个输入是%v1–需要保留从变量到绑定值的映射查找, %v1是一个日志。
不同的数据结构,影响编写转换的方式。为什么需要 let 绑定?PL 是一个相当成熟的领域。
为什么可能需要绑定
let 绑定的一个关键用法,指定计算范围。下面的例子,没有使用 let 绑定。
试图决定应该在哪里评估 node 时,问题就出现了%1。特别是,虽然文本格式似乎建议应该%1,在if 范围之外评估节点,但 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 放在哪里。数据流表单为后面的passes,决定将评估点放在哪里。在优化的初始阶段,当发现方便时,使用程序的数据流形式,可能不是一个坏主意。Relay 中的许多优化,都是为了优化数据流程序编写的。
将 IR 降低到实际运行时程序时,需要精确计算的范围。使用子函数和闭包时,希望明确指定计算范围,应该发生在哪里。Let-binding 可用于后期执行特定优化,解决问题。
对 IR 转换的影响
大多数函数式编程语言,用 A 范式分析,不需要表达式DAG。
Relay同时支持数据流形式和 let 绑定。让框架开发人员熟悉的表示。如何编写 pass,有一些影响:
• 如果来自数据流背景,要处理 let,将 var 映射到表达式,遇到 var 时,执行查找。需要一个从表达式到转换表达式的映射。有效地删除程序中的所有lets。
• 如果来自 PL 背景,喜欢 A 范式,将提供 A 范式pass的数据流。
• 对于 PL 人员,当实现某些东西(例如,数据流到 ANF 的转换)时,表达式可以是 DAG,应该使用 Map<Expr, Result> 访问表达式,计算一次转换后的结果,结果表达式保持公共结构。
有一些额外的高级概念,例如符号形状推理,未涵盖的多态函数。
参考链接:
https://tvm.apache.org/docs/dev/relay_intro.html