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
substract
LLVM 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示例:
在Parse的时候,由于还没有程序的全部信息,所以只可做局部的优化,如窥孔优化(例:常量折叠,Identity-function)。通过设计合适的类及继承体系,可以使用简单的算法实现peephole优化:
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)
<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的需求和挑战
将上述的表达式表示为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代码类似如下的伪代码:
函数式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
在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
@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
自由变量和闭包
自由变量(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
- [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
了解完MindSpore的关键技术是不是很心动呢!赶紧【点击链接】并【立即报名】,即可在 ModelArts 平台学习到一个经典案例掌握基于MindSpore的深度学习!
- 点赞
- 收藏
- 关注作者
评论(0)