​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​摘要：本文重点分析一下 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 asubstract

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 示例：

在这个示例中，方框为图的节点，表示 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 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_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 为：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 为：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 为：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 为：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]http://matt.might.net/articles

[6]《Don't Unroll Adjoint: Differentiating SSA-Form Programs》

[7] https://mindspore.cn/doc/note/z

点击关注，第一时间了解华为云新鲜技术~