MindSpore技术专栏 | AI框架中图层IR的分析
本文是MindSpore首席架构师金学峰的知乎专栏『AI框架分析』的第二篇，首先向大家介绍下IR是什么？
IR（Intermediate Representation即中间表示）是程序编译过程中，源代码与目标代码之间翻译的中介，IR的设计对编译器来说非常关键，好的IR要考虑从源代码到目标代码编译的完备性、编译优化的易用性和性能。而AI框架本质的作用又是什么呢？AI框架本质的作用在于把一个把用户的模型表达翻译到可执行的代码，然后进行高效执行（训练和推理），其中从用户的模型表达（例如深度神经网络）到最后可执行的代码就是个编译器的行为，这个编译器也有一个IR，它的设计对AI框架的完备性/灵活性/易用性/性能都起到至关重要的作用。
本文重点分析一下AI框架对IR有什么特殊的需求、业界有什么样的方案以及MindSpore的一些思考。首先带大家了解下通用的编译器IR的分类以及各自特点。
业界IR的介绍
Linear IR（线性IR）：
类似某些抽象机的伪代码，对应的算法通过迭代来遍历简单的线性操作序列
Hybrid IR（混合IR）：
结合了图IR和线性IR的要素。一种常见的混合IR使用底层的线性IR来表示无循环代码块，使用图IR来表示这些块之间的控制流
将编译过程的知识/信息保存在图中，对应的算法通过对图中的对象（节点、边、列表和树）操作来描述
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示例：
.png)
在Parse的时候，由于还没有程序的全部信息，所以只可做局部的优化，如窥孔优化（例：常量折叠，Identity-function）。通过设计合适的类及继承体系，可以使用简单的算法实现peephole优化：
.png)
def foo(x):
return x+1转换为CPS形式，k就是一个continuation：
def foo(x,k):
k(x+1)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)def h(q): i = n - 1
c(q) -> prodprimes(i, c)
i = n - 1
prodprimes(i, h).png)
<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的作用
业界在图层IR上的一些流派
.png)
将上述的表达式表示为SSA后，并插入J及计算AD，可以得到如下图表示的伪SSA代码：
上图中的 %6 这里节点称为“alpha node”，对应的是Primal中的节点%6，也就是上面一排的B3，“/”operation的反向函数
对于CFG间的控制流，需要对控制流进行反向分析，并在Primal CFG中插入适当的哑phi节点来记录和回放控制流。例如这一段计算power的代码：
.png)
对应的 Primal CFG中，插入了 %1 phi节点作为哑phi节点来记录控制流。然后在AD CFG中使用此 %1 来进行控制（%1记录通过入栈控制流，然后在AD CFG中通过出栈来回放控制流）
.png)
通过后续的代码优化，AD的Power代码类似如下的伪代码：
.png)
函数式IR
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, 5*5, 5) -> body_pow -> header_pow(0, 5*5*5, 5) -> after_pow() (此时return 5*5*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的设计思考
<ANode> ::= <ValueNode> | <ParameterNode>
<ParameterNode> ::= Parameter
<ValueNode> ::= Scalar | Named | Tensor | Type | Shape
| Primitive | MetaFuncGraph | FuncGraph
<CNode> ::= (<AnfNode> …)
<AnfNode> ::= <CNode> | <ANode>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
.png)
在ANF中每个表达式都用let表达式绑定为一个变量，通过对变量的引用来表示对表达式输出的依赖，而在MindIR中每个表达式都绑定为一个节点，通过节点与节点之间的有向边表示依赖关系
@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
.png)
@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
.png)
自由变量和闭包
自由变量(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
.png)
[1]《Engineering a Compiler》Second Edition，Chapter 5. Intermediate Representation
[2]《Combining Analyses, Combining Optimizations》
- 点赞
- 收藏
- 关注作者
评论（0）