AI框架中圖層IR的分析

摘要：本文重點分析一下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
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示例：

在這個示例中，方框為圖的節點，表示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, 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的設計思考

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

 

點擊關注，第一時間了解華為雲新鮮技術~