摘要：本文重点剖析一下 AI 框架对 IR 有什么非凡的需要、业界有什么样的计划以及 MindSpore 的一些思考。
本文分享自华为云社区《MindSpore 技术专栏 | AI 框架中图层 IR 的剖析》，原文作者：元气满满的少女月。
IR（Intermediate Representation 即两头示意）是程序编译过程中，源代码与指标代码之间翻译的中介，IR 的设计对编译器来说十分要害，好的 IR 要思考从源代码到指标代码编译的齐备性、编译优化的易用性和性能。而 AI 框架实质的作用又是什么呢？AI 框架实质的作用在于把一个把用户的模型表白翻译到可执行的代码，而后进行高效执行（训练和推理），其中从用户的模型表白（例如深度神经网络）到最初可执行的代码就是个编译器的行为，这个编译器也有一个 IR，它的设计对 AI 框架的齐备性 / 灵活性 / 易用性 / 性能都起到至关重要的作用。
本文重点剖析一下 AI 框架对 IR 有什么非凡的需要、业界有什么样的计划以及 MindSpore 的一些思考。首先带大家理解下通用的编译器 IR 的分类以及各自特点。
业界 IR 的介绍
一、IR 依据其组织构造[1]，能够分为：Linear IR（线性 IR）、Graphical IR（图 IR）、Hybrid IR（混合 IR），其中
Linear IR（线性 IR）：
相似某些抽象机的伪代码，对应的算法通过迭代来遍历简略的线性操作序列
Hybrid IR（混合 IR）：
联合了图 IR 和线性 IR 的因素。一种常见的混合 IR 应用底层的线性 IR 来示意无循环代码块，应用图 IR 来示意这些块之间的控制流
Graphical IR（图 IR）：
将编译过程的常识 / 信息保留在图中，对应的算法通过对图中的对象（节点、边、列表和树）操作来形容
线性 IR 的一个例子是堆栈机代码（Stack-Machine Code），它是一种单地址代码，假设操作数存在一个栈中。大多数操作从栈取得操作数，并将其后果推入栈中。例如：表达式 b-a* 3 对应的堆栈机代码如下：
push 3
push a
multiply
push a
substractLLVM IR 是一个典型的混合 IR，它蕴含了两个档次（CFG+BB）：
顶层是控制流图（Control Flow Graph，简写为 CFG），来示意基本块（Basic Block，简写为 BB）间的控制流。CFG 的每个节点（Node）为一个基本块，基本块 b1 和 b2 之间有一条边（Edge）：b1->b2，如果控制流可能从基本块 b1 的最初一条指令流向基本块 b2 的第一条指令
底层是基本块，在基本块中，每条指令是以 SSA（Static Single Assignment）模式出现，这些指令形成一个指令线性列表
Sea of Nodes IR（by Cliff Click）是一种典型的图 IR[2]，在这种 IR 中，简化了 CFG 图中 BB+SSA 指令的两层构造，去除了 BB，剩下只蕴含指令的一层构造。它通过引入了非凡的 REGION、IF、PROJECTION 指令，将 BB 块中的全序指令放松为显式的数据依赖和管制依赖，并且对管制依赖和数据依赖采纳雷同的示意形式和解决形式，这样就简化了 IR 的剖析和变换。如下为一个简略的 IR 示例：
在这个示例中，方框为图的节点，示意 SSA 指令，箭头为图的边；实心箭头示意管制依赖；空心箭头示意数据依赖。从这个示例中能够看到此 IR 中显式的蕴含了 use-def 依赖，不须要进行额定的计算。
基于此 IR 中显式的 use-def 信息，能够不便的实现两类优化：Parse time 优化（Pessimistic）全局优化(Optimistic)
在 Parse 的时候，因为还没有程序的全副信息，所以只可做部分的优化，如窥孔优化（例：常量折叠，Identity-function）。通过设计适合的类及继承体系，能够应用简略的算法实现 peephole 优化：
对于全局优化，比方 Sparse Conditional Constant Propagation（SCCP），也能够很简略的实现；首先是基于图中显式的 use-def 计算出 def-use chains，而后能够很容易的实现 SCCPSea of Nodes IR 提供了一种十分重要的思路：将依赖信息显式的示意在图 IR 中。FIRM IR 中连续了这个思路
二、从罕用编程语言的角度来剖析 IR，咱们又能够看到 IR 的模式分为了两个不同的营垒：一类是命令式编程语言的编译器 IR，另外一类是函数编程语言的编译器 IR 命令式编程语言的编译器 IR 以 SSA 为根本的组成模式，这里就不在赘述了，上面重点介绍一下函数式编程语言的 IR，在函数式编程语言的 IR 中，CPS 或者 ANF 是其根本的组成模式 1. Continuation-passing style（CPS）直译为：间断传递格调 CPS 示意这样一种模式：一个函数 f 除了它本身的参数外，总是有一个额定的参数 continuationcontinuation 也是一个函数，当 f 实现了本人的返回值计算之后，不是返回，而是将此返回值作为 continuation 的参数，调用 continuation。所以 CPS 模式的函数从模式上看它不会 return，当它要 return 的时候会将所有的参数传递给 continuation，让 continuation 持续去执行。比方：
def foo(x):
return x+1转换为 CPS 模式，k 就是一个 continuation：
def foo(x,k):
k(x+1)直观上看，函数不“return”，而是“continue”CPS 的长处是让如下的信息显式化：过程返回（调用一个 continuation），两头值（具备显式的名称），求值程序，尾调用（采纳雷同的 continuation 调用一个过程）比方如下的一段 python 代码，求小于 n 的所有素数的积。
def prodprimes(n):
if n == 1:
return 1
if isprime(n):
return n * prodprimes(n - 1)
return prodprimes(n - 1)当采纳 CPS 模式示意时：
def prodprimes(n, c):
def k(b):
if b == True:
m = n - 1
def j(p):
a = n * p
c(a)
prodprimes(m, j)
else:
def h(q):
c(q)
i = n - 1
prodprimes(i, h)
if n == 1:
c(1)
else:
isprime(n, k)从下面的代码中能够看到，“过程返回”都被调用 c、j、k、h 等 continuation 代替；两头值 a、b、m、i 都被给予了变量名称。CPS 模式非常适合编译器进行剖析和变换，比方 tail-recursion elimination 变换：如果函数 f 的最初是调用函数 g，那么函数 g 的 continuation 就不须要是在 f 内生成的一个 continuation，而能够被替换为传递给 f 的 continuation。下面的例子的原始代码中，“return prodprimes(n – 1)”语句就是一个尾递归在 CPS 模式中，能够很分明的看到 h(q）的定义其实就等于 c(q)，所以能够说 h 等于 c，于是能够进行如下的变换[3]：
def h(q): i = n - 1
c(q) -> prodprimes(i, c)
i = n - 1
prodprimes(i, h)尽管 CPS 十分统一和弱小，然而它的一个很大问题是难以浏览。所以呈现了 A -norm Form（ANF）模式 2. ANF 模式间接对 Direct Style 的源码进行转换[4]，不须要通过 CPS 模式
ANF 模式将表达式划分为两类：原子表达式和复合表达式。
原子表达式示意一个常数值或一个变量或一个原语或一个匿名函数复合表达式由多个原子表达式复合组成，能够看成是一个匿名函数或原语函数调用，组合的第一个输出是调用的函数，其余输出是调用的参数。一个复合表达式要么被 let-bound 到一个变量，要么只能呈现在最初的地位能够看到，ANF 模式通过 let-bound，显式表白了两头值和控制流及求值程序它的文法定义如下[5]
<aexp> ::= NUMBER | STRING | VAR | BOOLEAN | PRIMOP
| (lambda (VAR …) <exp>)
<cexp> ::= (<aexp> <aexp> …)
| (if <aexp> <exp> <exp>)
<exp> ::= (let ([VAR <cexp>]) <exp>) | <cexp> | <aexp>例如下面的 prodprimes 函数，如果用上述的文法示意，应该为：
(define prodprimes
(lambda (n)
(let (a (= n 1))
(if a 1 (let (b isprime(n))
(if b (let (m (- n 1))
(let (p (prodprimes m))
(* n p)))
(let (s (- n 1))
(prodprimes m))
))))))这段 ANF 模式表白，如果翻译为 python，应该相似于：
def prodprimes(n):
r = n == 1
if r:
return 1
b = isprime(n)
if b:
m = n - 1
p = prodprimes(m)
return n * p
s = n - 1
return prodprimes(s)通过这段代码，也能够看出，ANF 模式比 CPS 模式简略易懂
AI 框架中图层 IR 的作用
当初支流的 AI 框架都有图层 IR，好的图层 IR 有利于 AI 模型的编译优化和执行，是 AI 框架进行高效训练和推理的根底从训练的角度看，目前业界的 AI 框架有三种执行模式：Eager 执行模式、图执行模式和 Staging(混合)执行模式，其中高性能模式下（Graph 执行模式和 Staging 执行模式）都要基于图层 IR：Eager 执行模式个别是利用宿主语言（当初次要是 Python）的个性进行解释执行，外面应用了重载和 Tape 的一些技巧。
Graph 执行模式次要是拿到 AI 模型的图构造，而后进行编译优化和执行，这里的编译优化和执行就要基于图 IR，当初有三种办法拿到 AI 模型的图构造：第一种是程序员应用 API 构图（TF1.x 版本等）第二种是 Tracing JIT（JAX 带来的潮流，当初 TF2.0/Pytorch 等都反对）即把用户的模型脚本模仿跑一下，拿到正向的执行序列，而后基于这个序列进行构图益处是与 Eagle 模式比拟容易匹配，实现简略毛病是控制流的转换比拟麻烦、执行序列如果与算子执行后果相干的话不好实现、不容易解决副作用所以 TF 的 AutoGraph 还须要联合 AST 剖析解决控制流转换的问题第三种是 AST JIT（Pytorch 的 TorchScript）基于 Python 的 AST 进行构图，长处是转换的性能能够比拟全面，包含控制流等，毛病是实现简单，许多 Python 动静个性实现起来工作量大
Staging 执行模式相似在 Eager 模式中，通过 Python 修饰符，对局部子图进行编译执行减速（应用 Tracing JIT 或者 AST JIT），也会用到图 IR。
从推理的角度看，AI 框架生成最终的推理模型时须要进行大量的编译优化，如量化、剪枝等，个别都在图层 IR 上进行，同时最终的推理模型格局也是间接或者间接应用到图层 IRAI 框架图层 IR 的需要和挑战与其余通用的 IR 相比，AI 框架的图层 IR 有一些比拟非凡的需要和挑战：
张量表白：AI 的模型次要解决的是张量数据，这个与一般的利用差异是比拟大的，不过减少张量数据类型对编译器的 IR 来说并不是件艰难的事件。
主动微分：可微分是 AI 模型开发与个别利用开发区别最大的中央，古代的 AI 框架都会提供主动微分的性能，挑战在于实现的简洁性、性能以及将来高阶微分的扩大能力
JIT 能力：无论是图模式还是 Staging 模式，从算法工程师角度看，因为没有显示编译的步骤，都能够认为是 JIT 形式。对于 JIT 来说，编译性能是一个次要挑战
隐式并行：从开发者来说，有两种并行形式一种是是显式并行，开发者明确通知零碎哪里并行，比方显示启动多线程 / 增加
并行修饰符：还有一种形式是隐式并行，通过编译器来剖析依赖，主动实现并行一般而言，传统的 CFG+BB 的编译器，因为程序剖析应用全序剖析，不便做显式并行；函数式的编译器实践上易于数据依赖剖析，不便进行隐式并行优化。乏味的是，在深度学习场景中，Kernel 执行占了大部分开销，在运行时实现异步并发的模式也能够显著晋升整体性能，隐式并行的作用绝对会被弱化，然而想要实现极致性能，隐式并行还是有作用的
Loop 优化：AI 的计算波及大量的 Tensor 运算，对编译器来说就是 Loop 优化（张量—> 标量—> 向量化），不过这个挑战次要还是在算子层的 IR 上当然，图层 IR 也是是一种编译器 IR，应该具备通用性，包含类型零碎、控制流和数据流剖析、副作用打消等根本的性能
业界在图层 IR 上的一些流派
计算图的 IR：一种以 DAG 为核心的实现形式，许多晚期的框架都是应用了这种计划计算图的 IR 的设计比拟天然，计算图次要由边和节点组成，节点个别用来表白算子、变量、常量等等；边对应于 Tensors，实际上表白了一种数据依赖关系。前面的主动微分和优化都是基于这个 DAG 进行这个计划的长处是简略直观、优化时的性能开销小不足之处是计算图 IR 不算是真正形式化的编译器 IR，在类型零碎、简单逻辑的反对（比方递归）、副作用解决、控制流和数据流剖析方面反对不残缺
CFG+BB：基于传统编译器的 IR 来做图层 IR，比方 TorchScript、Julia 等如何实现主动微分？咱们以 Julia Zygote 为例[6]：对于 BB 块内的一般代码（非 phi，非 branch），借助链式法则，能够依照反向的程序生成 AD 代码
将上述的表达式示意为 SSA 后，并插入 J 及计算 AD，能够失去如下图示意的伪 SSA 代码：
上图中的 %6 这里节点称为“alpha node”，对应的是 Primal 中的节点 %6，也就是下面一排的 B3，“/”operation 的反向函数
对于 CFG 间的控制流，须要对控制流进行反向剖析，并在 Primal CFG 中插入适当的哑 phi 节点来记录和回放控制流。例如这一段计算 power 的代码：
对应的 Primal CFG 中，插入了 %1 phi 节点作为哑 phi 节点来记录控制流。而后在 AD CFG 中应用此 %1 来进行管制（%1 记录通过入栈控制流，而后在 AD CFG 中通过出栈来回放控制流）
通过后续的代码优化，AD 的 Power 代码相似如下的伪代码：
能够看出，CFG+BB 的主动微分最终是通过迭代的形式来实现的，带 Scope 的 SSA 模式须要解决边界传递的问题对主动微分还是会带来一些解决上的麻烦
如何做图优化？转化成 use-def、def-use 的模式进行优化
如何做并行优化？因为 CFG+BB 是全序的形式，须要转换成 use-def，并联合副作用信息进行剖析
应用 CFG+BB 计划的益处是性能齐备、计划成熟、重用性高，不过 CFG+BB 的模式对主动微分 / 图优化 / 并行优化来说，都要进行肯定的转换工作，并不是那么直观和高效
函数式 IR
应用函数式的 IR 来做图层 IR，典型的如 Relay、Myia 等，如何实现主动微分？对于非控制流，计算 AD 的办法和上述的 BB 块内计算 AD 的办法雷同。对于控制流，函数式 IR 采纳了不同的解决形式，将迭代转换为递归，并且通过 switch 函数来进行分支的抉择。例如上述雷同的 pow()函数：
def pow(x, n):
return header_pow(n, 1, x)
def header_pow(phi_n, phi_r, x):
def body_pow():
phi_n_1 = phi_n - 1
phi_r_1 = phi_r * x
return header_pow(phi_n_1, phi_r_1, x)
def after_pow():
return phi_r
f = switch(phi_n > 0, header_pow, after_pow)
f()以 pow(5,3) 为例，其递归调用过程如下：
pow(5, 3) -> header_pow(3, 1, 5) -> body_pow() -> header_pow(2, 5, 5) -> body_pow() -> header_pow(1, 55, 5) -> body_pow -> header_pow(0, 555, 5) -> after_pow() (此时 return 55*5)
能够看到，这里的递归调用的调用和返回别离就对应了上述 CFG+BB 的控制流 phi 节点入栈和出栈操作
因为 AD 过程就是对函数进行变换的过程，所以 AD 后的图也是递归调用的构造，因而不须要相似 CFG+BB 的控制流 phi 节点入栈和出栈操作，递归调用过程人造的就代替了入栈和出栈的过程
对 x 求导数
def x_grad_pow(x, n):
phi_n = n
phi_r = 1
return x_bprop_header_pow(phi_n, phi_r, x, 1)
def x_bprop_header_pow(phi_n, phi_r, x, sens):
def env_x_bprop_body_pow():
%3 = x_bprop_header_pow(phi_n – 1, phi_r * phi_x, x, 1)
%4 = phi_r_bprop_header_pow(phi_n – 1, phi_r * phi_x, x, 1)
%5 = %4 * phi_r
return %3 + %5
def env_x_bprop_after_pow():
return 0
f = switch(phi_n > 0, env_x_bprop_body_pow, env_x_bprop_after_pow)
r = switch(phi_n > 0, f(), 0)
return r
def phi_r_bprop_header_pow(phi_n, phi_r, x, sens):
def env_phi_r_bprop_body_pow():
%3 = phi_r_bprop_header_pow(phi_n - 1, phi_r * x, x, 1)
%4 = %3 * x
return %4
def env_phi_r_bprop_after_pow():
return 1
if phi_n > 0:
%5 = env_phi_r_bprop_body_pow()
else:
%5 = env_phi_r_bprop_after_pow()
return %5函数式 IR 的益处是对主动微分敌对，比拟适宜做并行剖析，不过挑战在于函数 IR 的副作用打消以及函数式 IR 在执行态的性能（含有递归对执行不敌对）
Mindspore 的设计思考
MindSpore 的图层 IR 叫做 MindIR，MindIR 抉择的技术路线是采纳 Functional Graph IR（参考了 Sea of Nodes、Thorin、Myia 等），具备如下特色：
Functional 以更天然的主动微分实现形式和更不便的隐式并行剖析能力：函数作为一等公民，反对高阶函数，包含控制流也通过非凡的函数来实现，能够以对立的模式来实现微分函数以无副作用的形式实现，与命令式语言相比，可简化剖析和实现更多的优化原生反对闭包，一方面能够不便的表白用户源代码中的闭包示意，另外也能够天然的反对主动微分算法中在反向函数中要拜访原始函数的两头后果的要求：反向函数拜访两头后果，并且作为一个闭包返回应用基于数据依赖的偏序剖析，这样能够便于乱序或者并行执行
Graph based 以更适宜 JIT 的疾速优化能力：采纳相似 Sea of Nodes IR 的只有一层的示意形式，控制流和数据流合一，更适宜 JIT 优化
ANF 模式：和 Thorin 相似，都采纳 Graph IR，都打消了 Scope。然而没有采纳 Thorin IR 的 CPS 模式，而是表达能力相似，更直观也更易查看的 ANF 模式 MindIR 心愿通过 Functional 的形式更不便的实现主动微分和隐式并行剖析，Graph Based 形式把控制流和数据流合一反对更高效的 JIT 优化。一、MindIR 的详解[7]MindIR 文法继承于 ANF，其定义如下所示：
<ANode> ::= <ValueNode> | <ParameterNode>
<ParameterNode> ::= Parameter
<ValueNode> ::= Scalar | Named | Tensor | Type | Shape
| Primitive | MetaFuncGraph | FuncGraph
<CNode> ::= (<AnfNode> …)
<AnfNode> ::= <CNode> | <ANode>MindIR 中的 ANode 对应于 ANF 的原子表达式，ANode 有两个子类别离为 ValueNode 和 ParameterNodeValueNode 示意常数节点可承载一个常数值（标量、符号、张量、类型、维度等），也能够是一个原语函数（Primitive）或一个元函数（MetaFuncGraph）或一个一般函数（FuncGraph），因为在函数式编程中函数定义自身也是一个值，ParameterNode 是参数节点示意函数的形参 MindIR 中 CNode 对应于 ANF 的复合表达式，示意一次函数调用在 MindSpore 主动微分时，会计算 ParameterNode 和 CNode 的梯度奉献，并返回最终 ParameterNode 的梯度，而不计算 ValueNode 的梯度
上面以一段程序作为示例，比照了解 MindIR
def func(x, y):
return x / y
@ms_function
def test_f(x, y):
a = x - 1
b = a + y
c = b * func(a, b)
return c这段 Python 代码对应的 ANF 表白为：
lambda (x, y)
let a = x - 1 in
let b = a + y in
let func = lambda (x, y)
let ret = x / y in
ret end in
let %1 = func(a, b) in
let c = b * %1 in
c end对应的 MindIR 为：https://w.url.cn/s/Ansh1KW
在 MindIR 中，一个函数图（FuncGraph）示意一个一般函数的定义，函数图个别由 ParameterNode、ValueNode 和 CNode 组成有向无环图，能够清晰地表白出从参数到返回值的计算过程在上图中能够看出，python 代码中两个函数 test_f 和 func 转换成了两个函数图，其参数 x 和 y 转换为函数图的 ParameterNode，每一个表达式转换为一个 CNode。CNode 的第一个输出链接着调用的函数，例如图中的 add、func、return 值得注意的是这些节点均是 ValueNode，因为它们被了解为常数函数值。CNode 的其余输出链接这调用的参数，参数值能够来自于 ParameterNode、ValueNode 和其余 CNode。
在 ANF 中每个表达式都用 let 表达式绑定为一个变量，通过对变量的援用来示意对表达式输入的依赖，而在 MindIR 中每个表达式都绑定为一个节点，通过节点与节点之间的有向边示意依赖关系
函数式语义
MindIR 较传统计算图的一个重要个性是不仅能够表白算子之间的数据依赖，还能够表白丰盛的函数式语义
高阶函数
在 MindIR 中，函数的定义是由一个子图来定义，但其自身能够是一个被传递的值，作为其余高阶函数的输出或输入。例如上面一个简略的示例中，函数 f 作为参数传入了函数 g，因而函数 g 是一个接管函数输出的高阶函数，函数 f 真正的调用点是在函数 g 外部
@ms_function
def hof(x):
def f(x):
return x + 3
def g(function, x):
return function(x) * function(x)
res = g(f, x)
return res对应的 MindIR 为：https://w.url.cn/s/A8vb8X3
在理论网络训练脚本中，主动求导泛函 GradOperation 和优化器中罕用到的 Partial 和 HyperMap 都是典型的高阶函数。高阶语义极大地晋升了 MindSpore 表白的灵活性和简洁性
控制流
控制流在 MindIR 中是以高阶函数抉择调用的模式表白。这样的模式把控制流转换为高阶函数的数据流，从而使得主动微分算法更加弱小。不仅能够反对数据流的主动微分，还能够反对条件跳转、循环和递归等控制流的主动微分。上面以一个简略的斐波那契用例来演示阐明
@ms_function
def fibonacci(n):
if(n < 1):
return 0
elif(n == 1):
return 1
else:
return fibonacci(n-1) + fibonacci(n-2)对应的 MindIR 为：https://w.url.cn/s/AUiE9Mc
其中 fibonacci 是顶层函数图，在顶层中有两个函数图被 switch 抉择调用✓fibonacci 是第一个 if 的 True 分支，✗fibonacci 是第一个 if 的 False 分支。在✗fibonacci 中被调用的✓✗fibonacci 是 elif 的 True 分支，✗✗fibonacci 是 elif 的 False 分支。
这里须要了解的要害是在 MindIR 中，条件跳转和递归是以高阶控制流的模式表白的例如，✓fibonacci 和✗fibonacci 是作为 switch 算子的参数传入，switch 依据条件参数抉择哪一个函数作为返回值因而，switch 是把输出的函数当成一般的值做了一个二元抉择操作，并没有调用，而真正的函数调用是在紧随 switch 后的 CNode 上实现
自在变量和闭包
自在变量 (free variable) 是指在代码块中援用作用域环境中的变量而非局部变量
闭包（closure）是一种编程语言个性，它指的是代码块和作用域环境的联合
在 MindIR 中，代码块是以函数图出现的，而作用域环境能够了解为该函数被调用时的上下文环境，自在变量的捕捉形式是值拷贝而非援用。
一个典型的闭包用例如下：
@ms_function
def func_outer(a, b):
def func_inner(c):
return a + b + c
return func_inner
@ms_function
def ms_closure():
closure = func_outer(1, 2)
out1 = closure(1)
out2 = closure(2)
return out1, out2对应的 MindIR 为：https://w.url.cn/s/AsUMXTS
在例子中，a 和 b 是自在变量，因为 func_inner 中变量 a 和 b 是援用的其父图 func_outer 中定义的参数。变量 closure 是一个闭包，它是函数 func_inner 与其上下文 func_outer(1, 2)的联合。因而，out1 的后果是 4，因为其等价于 1 +2+1，out2 的后果是 5，因为其等价于 1 +2+2
参考文献
[1]《Engineering a Compiler》Second Edition，Chapter 5. Intermediate Representation
[2]《Combining Analyses, Combining Optimizations》
[3]《COMPILING WITH CONTINUATIONS》第一章
[4]《Functional programming languages Part V: functional intermediate representations》
[5] matt.might.net/articles
[6]《Don’t Unroll Adjoint: Differentiating SSA-Form Programs》
[7] mindspore.cn/doc/note/z
点击关注，第一工夫理解华为云陈腐技术~