AI框架中圖層IR的分析

摘要：本文重點分析一下AI框架對IR有什麼特殊的需求、業界有什麼樣的方案以及MindSpore的一些思考。

本文分享自華為雲社群《》，原文作者：元氣滿滿的少女月 。

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為：

在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為：

在實際網路訓練指令碼中，自動求導泛函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為：

其中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為：

在例子中，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