AI框架中圖層IR的分析

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

本文分享自華爲雲社區《MindSpore技術專欄 | AI框架中圖層IR的分析》，原文做者：元氣滿滿的少女月 。

IR（Intermediate Representation即中間表示）是程序編譯過程當中，源代碼與目標代碼之間翻譯的中介，IR的設計對編譯器來講很是關鍵，好的IR要考慮從源代碼到目標代碼編譯的完備性、編譯優化的易用性和性能。而AI框架本質的做用又是什麼呢？AI框架本質的做用在於把一個把用戶的模型表達翻譯到可執行的代碼，而後進行高效執行（訓練和推理），其中從用戶的模型表達（例如深度神經網絡）到最後可執行的代碼就是個編譯器的行爲，這個編譯器也有一個IR，它的設計對AI框架的完備性/靈活性/易用性/性能都起到相當重要的做用。

本文重點分析一下AI框架對IR有什麼特殊的需求、業界有什麼樣的方案以及MindSpore的一些思考。首先帶你們瞭解下通用的編譯器IR的分類以及各自特色。

業界IR的介紹

1、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中延續了這個思路

2、從經常使用編程語言的角度來分析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, 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優化。1、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

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