摘要：本文重点剖析一下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 3push amultiplypush asubstractLLVM 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 - 1prodprimes(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 - 1return 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 rdef 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_functiondef 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_functiondef 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_functiondef 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_functiondef func_outer(a, b): def func_inner(c): return a + b + c return func_inner@ms_functiondef 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
点击关注，第一工夫理解华为云陈腐技术~