轉(zhuǎn)載自：知乎
作者：金雪鋒

AI框架圖層IR的分析

前段時(shí)間一直忙于HC大會(huì)和MindSpore1.0的版本準(zhǔn)備，國(guó)慶期間終于有點(diǎn)時(shí)間來交作業(yè)了，本文是AI框架分析專欄的第二篇

背景

IR-Intermediate Representation（中間表示）是程序編譯過程中，源代碼與目標(biāo)代碼之間翻譯的中介，IR的設(shè)計(jì)對(duì)編譯器來說非常關(guān)鍵，好的IR要考慮從源代碼到目標(biāo)代碼編譯的完備性、編譯優(yōu)化的易用性和性能。

AI框架本質(zhì)的作用在于把一個(gè)把用戶的模型表達(dá)翻譯到可執(zhí)行的代碼，然后進(jìn)行高效執(zhí)行（訓(xùn)練和推理），其中從用戶的模型表達(dá)（例如深度神經(jīng)網(wǎng)絡(luò)）到最后可執(zhí)行的代碼就是個(gè)編譯器的行為，這個(gè)編譯器也有一個(gè)IR，它的設(shè)計(jì)對(duì)AI框架的完備性/靈活性/易用性/性能都起到至關(guān)重要的作用；本文重點(diǎn)分析一下AI框架對(duì)IR有什么特殊的需求、業(yè)界有什么樣的方案以及MindSpore的一些思考。

在開始之前，我想先介紹一些背景知識(shí)—通用的編譯器IR的分類以及各自特點(diǎn)

業(yè)界IR的介紹

參考[1]，IR根據(jù)其組織結(jié)構(gòu)，可以分為：

• Linear IR（線性IR）：類似某些抽象機(jī)的偽代碼，對(duì)應(yīng)的算法通過迭代來遍歷簡(jiǎn)單的線性操作序列。
• Graphical IR（圖IR）：將編譯過程的知識(shí)/信息保存在圖中，對(duì)應(yīng)的算法通過對(duì)圖中的對(duì)象（節(jié)點(diǎn)、邊、列表和樹）操作來描述。
• Hybrid IR（混合IR）：結(jié)合了圖IR和線性IR的要素。一種常見的混合IR使用底層的線性IR來表示無循環(huán)代碼塊，使用圖IR來表示這些塊之間的控制流。

線性IR的一個(gè)例子是堆棧機(jī)代碼（Stack-Machine Code），它是一種單地址代碼，假定操作數(shù)存在一個(gè)棧中。大多數(shù)操作從棧獲得操作數(shù)，并將其結(jié)果推入棧中。例如：表達(dá)式 b-a*3對(duì)應(yīng)的堆棧機(jī)代碼如下：

push 3
push a
multiply
push a
substract

LLVM IR是一個(gè)典型的混合IR，它包含了兩個(gè)層次（CFG+BB）：

頂層是控制流圖（Control Flow Graph，簡(jiǎn)寫為CFG），來表示基本塊（Basic Block，簡(jiǎn)寫為BB）間的控制流。CFG的每個(gè)節(jié)點(diǎn)（Node）為一個(gè)基本塊，基本塊b1和b2之間有一條邊（Edge）：b1->b2，如果控制流可能從基本塊b1的最后一條指令流向基本塊b2的第一條指令。

底層是基本塊，在基本塊中，每條指令是以SSA（Static Single Assignment）形式呈現(xiàn)，這些指令構(gòu)成一個(gè)指令線性列表。

Sea of Nodes IR（by Cliff Click）是一種典型的圖IR，參考[2]，在這種IR中，簡(jiǎn)化了CFG圖中BB+SSA指令的兩層結(jié)構(gòu)，去除了BB，剩下只包含指令的一層結(jié)構(gòu)。它通過引入了特殊的REGION、IF、PROJECTION指令，將BB塊中的全序指令放松為顯式的數(shù)據(jù)依賴和控制依賴，并且對(duì)控制依賴和數(shù)據(jù)依賴采用相同的表示方式和處理方式，這樣就簡(jiǎn)化了IR的分析和變換。如下為一個(gè)簡(jiǎn)單的IR示例：

在這個(gè)示例中，方框?yàn)閳D的節(jié)點(diǎn)，表示SSA指令，箭頭為圖的邊；實(shí)心箭頭表示控制依賴；空心箭頭表示數(shù)據(jù)依賴。從這個(gè)示例中可以看到此IR中顯式的包含了use-def依賴，不需要進(jìn)行額外的計(jì)算。

基于此IR中顯式的use-def信息，可以方便的實(shí)現(xiàn)兩類優(yōu)化：1、Parse time優(yōu)化（Pessimistic）；2、全局優(yōu)化(Optimistic)；

在Parse的時(shí)候，由于還沒有程序的全部信息，所以只可做局部的優(yōu)化，如窺孔優(yōu)化（例：常量折疊，Identity-function）。通過設(shè)計(jì)合適的類及繼承體系，可以使用簡(jiǎn)單的算法實(shí)現(xiàn)peephole優(yōu)化：

對(duì)于全局優(yōu)化，比如Sparse Conditional Constant Propagation（SCCP），也可以很簡(jiǎn)單的實(shí)現(xiàn)；首先是基于圖中顯式的use-def計(jì)算出def-use chains，然后可以很容易的實(shí)現(xiàn)SCCP。具體可以參考論文。

Sea of Nodes IR提供了一種非常重要的思路：將依賴信息顯式的表示在圖IR中。FIRM IR中延續(xù)了這個(gè)思路。

從常用編程語言的角度來分析IR，我們又可以看到IR的形式分為了兩個(gè)不同的陣營(yíng)，一類是命令式編程語言的編譯器IR，另外一類是函數(shù)編程語言的編譯器IR；命令式編程語言的編譯器IR以SSA為基本的組成形式，這里就不在贅述了，下面重點(diǎn)介紹一下函數(shù)式編程語言的IR，在函數(shù)式編程語言的IR中，CPS或者ANF是其基本的組成形式。

Continuation-passing style（CPS）直譯為：連續(xù)傳遞風(fēng)格。CPS 表示這樣一種形式：一個(gè)函數(shù) f 除了它自身的參數(shù)外，總是有一個(gè)額外的參數(shù)continuation。continuation也是一個(gè)函數(shù)，當(dāng)f完成了自己的返回值計(jì)算之后，不是返回，而是將此返回值作為continuation的參數(shù)，調(diào)用continuation。所以CPS形式的函數(shù)從形式上看它不會(huì)return，當(dāng)它要return 的時(shí)候會(huì)將所有的參數(shù)傳遞給continuation，讓continuation繼續(xù)去執(zhí)行。比如：

def foo(x):
return x+1

轉(zhuǎn)換為CPS形式，k就是一個(gè)continuation：

def foo(x,k):
k(x+1)

直觀上看，函數(shù)不“return”，而是“continue”;

CPS的優(yōu)點(diǎn)是讓如下的信息顯式化：過程返回（調(diào)用一個(gè)continuation），中間值（具有顯式的名稱），求值順序，尾調(diào)用（采用相同的continuation調(diào)用一個(gè)過程）。

比如如下的一段python代碼，求小于n的所有素?cái)?shù)的積。

def prodprimes(n):
    if n == 1:
        return 1
    if isprime(n):
        return n * prodprimes(n - 1)
return prodprimes(n - 1)

當(dāng)采用CPS形式表示時(shí)：

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)

從上面的代碼中可以看到，“過程返回”都被調(diào)用c、j、k、h等continuation代替；中間值a、b、m、i都被給予了變量名稱。

CPS形式非常適合編譯器進(jìn)行分析和變換，比如tail-recursion elimination變換：如果函數(shù)f的最后是調(diào)用函數(shù)g，那么函數(shù)g的continuation就不需要是在f內(nèi)生成的一個(gè)continuation，而可以被替換為傳遞給f的continuation。上面的例子的原始代碼中，“return prodprimes(n - 1)”語句就是一個(gè)尾遞歸。在CPS形式中，可以很清楚的看到h(q）的定義其實(shí)就等于c(q)，所以可以說h等于c，于是可以進(jìn)行如下的變換：

def h(q):                         i = n - 1
    c(q)            ->           prodprimes(i, c)
i = n - 1
prodprimes(i, h)

參考【3】

雖然CPS非常一致和強(qiáng)大，但是它的一個(gè)很大問題是難以閱讀。所以出現(xiàn)了A-norm Form（ANF）形式。ANF形式直接對(duì)Direct Style的源碼進(jìn)行轉(zhuǎn)換，不需要經(jīng)過CPS形式。參考[4]

ANF形式將表達(dá)式劃分為兩類：原子表達(dá)式和復(fù)合表達(dá)式，原子表達(dá)式表示一個(gè)常數(shù)值或一個(gè)變量或一個(gè)原語或一個(gè)匿名函數(shù)；復(fù)合表達(dá)式由多個(gè)原子表達(dá)式復(fù)合組成，可以看成是一個(gè)匿名函數(shù)或原語函數(shù)調(diào)用，組合的第一個(gè)輸入是調(diào)用的函數(shù)，其余輸入是調(diào)用的參數(shù)。一個(gè)復(fù)合表達(dá)式要么被let-bound到一個(gè)變量，要么只能出現(xiàn)在最后的位置。可以看到，ANF形式通過let-bound，顯式表達(dá)了中間值和控制流及求值順序。

它的文法定義如下，參考[5]：

 ::= NUMBER | STRING | VAR | BOOLEAN | PRIMOP
          |  (lambda (VAR …) )
 ::= ( …)
          |  (if )
 ::= (let ([VAR ]) ) |  | 

例如上面的prodprimes函數(shù)，如果用上述的文法表示，應(yīng)該為：

(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形式表達(dá)，如果翻譯為python：

應(yīng)該類似于：

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形式簡(jiǎn)單易懂。

AI框架中圖層IR的作用

現(xiàn)在主流的AI框架都有圖層IR，好的圖層IR有利于AI模型的編譯優(yōu)化和執(zhí)行，是AI框架進(jìn)行高效訓(xùn)練和推理的基礎(chǔ)。

從訓(xùn)練的角度看，目前業(yè)界的AI框架有三種執(zhí)行模式：Eager執(zhí)行模式、圖執(zhí)行模式和Staging(混合)執(zhí)行模式，其中高性能模式下（Graph執(zhí)行模式和Staging執(zhí)行模式）都要基于圖層IR：

• Eager執(zhí)行模式一般是利用宿主語言（現(xiàn)在主要是Python）的特性進(jìn)行解釋執(zhí)行，里面使用了重載和Tape的一些技巧。
• Graph執(zhí)行模式主要是拿到AI模型的圖結(jié)構(gòu)，然后進(jìn)行編譯優(yōu)化和執(zhí)行，這里的編譯優(yōu)化和執(zhí)行就要基于圖IR，額外提一下，現(xiàn)在有三種方法拿到AI模型的圖結(jié)構(gòu)，第一種是程序員使用API構(gòu)圖（TF1.x版本等）；第二種是Tracing JIT（JAX帶來的潮流，現(xiàn)在TF2.0/Pytorch等都支持），即把用戶的模型腳本模擬跑一下，拿到正向的執(zhí)行序列，然后基于這個(gè)序列進(jìn)行構(gòu)圖，好處是與Eagle模式比較容易匹配，實(shí)現(xiàn)簡(jiǎn)單，缺點(diǎn)是控制流的轉(zhuǎn)換比較麻煩、執(zhí)行序列如果與算子執(zhí)行結(jié)果相關(guān)的話不好實(shí)現(xiàn)、不容易處理副作用，所以TF的AutoGraph還需要結(jié)合AST分析解決控制流轉(zhuǎn)換的問題；第三種是AST JIT（Pytorch的TorchScript），基于Python的AST進(jìn)行構(gòu)圖，優(yōu)點(diǎn)是轉(zhuǎn)換的功能可以比較全面，包括控制流等，缺點(diǎn)是實(shí)現(xiàn)復(fù)雜，許多Python動(dòng)態(tài)特性實(shí)現(xiàn)起來工作量大。
• Staging執(zhí)行模式類似在Eager模式中，通過Python修飾符，對(duì)部分子圖進(jìn)行編譯執(zhí)行加速（使用Tracing JIT或者AST JIT），也會(huì)用到圖IR。

從推理的角度看，AI框架生成最終的推理模型時(shí)需要進(jìn)行大量的編譯優(yōu)化，如量化、剪枝等，一般都在圖層IR上進(jìn)行，同時(shí)最終的推理模型格式也是直接或者間接使用到圖層IR。

AI框架圖層IR的需求和挑戰(zhàn)

與其他通用的IR相比，AI框架的圖層IR有一些比較特殊的需求和挑戰(zhàn)：

• 張量表達(dá)

AI的模型主要處理的是張量數(shù)據(jù)，這個(gè)與普通的應(yīng)用差別是比較大的，不過增加張量數(shù)據(jù)類型對(duì)編譯器的IR來說并不是件困難的事情。

• 自動(dòng)微分

可微分是AI模型開發(fā)與一般應(yīng)用開發(fā)區(qū)別最大的地方，現(xiàn)代的AI框架都會(huì)提供自動(dòng)微分的功能，挑戰(zhàn)在于實(shí)現(xiàn)的簡(jiǎn)潔性、性能以及未來高階微分的擴(kuò)展能力。

• JIT能力

無論是圖模式還是Staging模式，從算法工程師角度看，由于沒有顯示編譯的步驟，都可以認(rèn)為是JIT方式。對(duì)于JIT來說，編譯性能是一個(gè)主要挑戰(zhàn)。

• 隱式并行

從開發(fā)者來說，有兩種并行方式，一種是是顯式并行，開發(fā)者明確告訴系統(tǒng)哪里并行，比如顯示啟動(dòng)多線程/添加并行修飾符；還有一種方式是隱式并行，通過編譯器來分析依賴，自動(dòng)實(shí)現(xiàn)并行。一般而言，傳統(tǒng)的CFG+BB的編譯器，由于程序分析使用全序分析，方便做顯式并行；函數(shù)式的編譯器理論上易于數(shù)據(jù)依賴分析，方便進(jìn)行隱式并行優(yōu)化。有趣的是，在深度學(xué)習(xí)場(chǎng)景中，Kernel執(zhí)行占了大部分開銷，在運(yùn)行時(shí)實(shí)現(xiàn)異步并發(fā)的模式也可以顯著提升整體性能，隱式并行的作用相對(duì)會(huì)被弱化，但是想要實(shí)現(xiàn)極致性能，隱式并行還是有作用的。

• Loop優(yōu)化

AI的計(jì)算涉及大量的Tensor運(yùn)算，對(duì)編譯器來說就是Loop優(yōu)化（張量—>標(biāo)量—>向量化），不過這個(gè)挑戰(zhàn)主要還是在算子層的IR上。

當(dāng)然，圖層IR也是是一種編譯器IR，應(yīng)該具備通用性，包括類型系統(tǒng)、控制流和數(shù)據(jù)流分析、副作用消除等基本的功能。

業(yè)界在圖層IR上的一些流派

• 計(jì)算圖的IR：一種以DAG為中心的實(shí)現(xiàn)方式，許多早期的框架都是使用了這種方案。他的設(shè)計(jì)比較自然，計(jì)算圖主要由邊和節(jié)點(diǎn)組成，節(jié)點(diǎn)一般用來表達(dá)算子、變量、常量等等；邊對(duì)應(yīng)于Tensors，實(shí)際上表達(dá)了一種數(shù)據(jù)依賴關(guān)系。后面的自動(dòng)微分和優(yōu)化都是基于這個(gè)DAG進(jìn)行，這一塊網(wǎng)上的資料很多，我就不贅述。

總結(jié)：這個(gè)方案的優(yōu)點(diǎn)是簡(jiǎn)單直觀、優(yōu)化時(shí)的性能開銷小；不足之處是計(jì)算圖IR不算是真正形式化的編譯器IR，在類型系統(tǒng)、復(fù)雜邏輯的支持（比如遞歸）、副作用處理、控制流和數(shù)據(jù)流分析方面支持不完整。

• CFG+BB：基于傳統(tǒng)編譯器的IR來做圖層IR，比如TorchScript、Julia等。

如何實(shí)現(xiàn)自動(dòng)微分？我們以Julia Zygote為例，參考[6]：

對(duì)于BB塊內(nèi)的普通代碼（非phi，非branch），借助鏈?zhǔn)椒▌t，可以按照反向的順序生成AD代碼。

將上述的表達(dá)式表示為SSA后，并插入J及計(jì)算AD，可以得到如下圖表示的偽SSA代碼：

上圖中的 %6 這里節(jié)點(diǎn)稱為“alpha node”，對(duì)應(yīng)的是Primal中的節(jié)點(diǎn)%6，也就是上面一排的B3，“/”operation的反向函數(shù)。

對(duì)于CFG間的控制流，需要對(duì)控制流進(jìn)行反向分析，并在Primal CFG中插入適當(dāng)?shù)膯hi節(jié)點(diǎn)來記錄和回放控制流。例如這一段計(jì)算power的代碼：

對(duì)應(yīng)的 Primal CFG中，插入了 %1 phi節(jié)點(diǎn)作為啞phi節(jié)點(diǎn)來記錄控制流。然后在AD CFG中使用此 %1 來進(jìn)行控制（%1記錄通過入棧控制流，然后在AD CFG中通過出棧來回放控制流）。

注：AD CFG中的 block #1中g(shù)oto #4，感覺應(yīng)該是goto #2；block #2中g(shù)oto #2，應(yīng)該是goto #1。

通過后續(xù)的代碼優(yōu)化，AD的Power代碼類似如下的偽代碼：

可以看出，CFG+BB的自動(dòng)微分最終是通過迭代的方式來實(shí)現(xiàn)的，帶Scope的SSA形式需要解決邊界傳遞的問題對(duì)自動(dòng)微分還是會(huì)帶來一些處理上的麻煩。

如何做圖優(yōu)化？

轉(zhuǎn)化成use-def、def-use的形式進(jìn)行優(yōu)化。

如何做并行優(yōu)化？

由于CFG+BB是全序的方式，需要轉(zhuǎn)換成use-def，并結(jié)合副作用信息進(jìn)行分析。

總結(jié)：使用CFG+BB方案的好處是功能完備、方案成熟、重用性高，不過CFG+BB的形式對(duì)自動(dòng)微分/圖優(yōu)化/并行優(yōu)化來說，都要進(jìn)行一定的轉(zhuǎn)換工作，并不是那么直觀和高效。

函數(shù)式IR：使用函數(shù)式的IR來做圖層IR，典型的如Relay、Myia等。

如何實(shí)現(xiàn)自動(dòng)微分？

對(duì)于非控制流，計(jì)算AD的方法和上述的BB塊內(nèi)計(jì)算AD的方法相同。對(duì)于控制流，函數(shù)式IR采用了不同的處理方式，將迭代轉(zhuǎn)換為遞歸，并且通過switch函數(shù)來進(jìn)行分支的選擇。例如上述相同的pow()函數(shù)：

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) 為例，其遞歸調(diào)用過程如下：

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() (此時(shí)return 5*5*5)

可以看到，這里的遞歸調(diào)用的調(diào)用和返回分別就對(duì)應(yīng)了上述CFG+BB的控制流phi節(jié)點(diǎn)入棧和出棧操作。

由于AD過程就是對(duì)函數(shù)進(jìn)行變換的過程，所以AD后的圖也是遞歸調(diào)用的結(jié)構(gòu)，因此不需要類似CFG+BB的控制流phi節(jié)點(diǎn)入棧和出棧操作，遞歸調(diào)用過程天然的就代替了入棧和出棧的過程。

/# 對(duì)x求導(dǎo)數(shù)

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

總結(jié)：函數(shù)式IR的好處是對(duì)自動(dòng)微分友好，比較適合做并行分析，不過挑戰(zhàn)在于函數(shù)IR的副作用消除以及函數(shù)式IR在執(zhí)行態(tài)的性能（含有遞歸對(duì)執(zhí)行不友好）。

MindSpore的設(shè)計(jì)思考

MindSpore的圖層IR叫做MindIR，MindIR選擇的技術(shù)路線是采用Functional Graph IR（參考了Sea of Nodes 、Thorin、Myia等），具有如下特征：

• Functional—更自然的自動(dòng)微分實(shí)現(xiàn)方式和更方便的隱式并行分析能力：函數(shù)作為一等公民，支持高階函數(shù)，包括控制流也通過特殊的函數(shù)來實(shí)現(xiàn)，可以以統(tǒng)一的形式來實(shí)現(xiàn)微分；函數(shù)以無副作用的方式實(shí)現(xiàn)，與命令式語言相比，可簡(jiǎn)化分析和實(shí)現(xiàn)更多的優(yōu)化；原生支持閉包，一方面可以方便的表達(dá)用戶源代碼中的閉包表示，另外也可以自然的支持自動(dòng)微分算法中在反向函數(shù)中要訪問原始函數(shù)的中間結(jié)果的要求：反向函數(shù)訪問中間結(jié)果，并且作為一個(gè)閉包返回；使用基于數(shù)據(jù)依賴的偏序分析，這樣可以便于亂序或者并行執(zhí)行。
• Graph based—更適合JIT的快速優(yōu)化能力：采用類似Sea of Nodes IR的只有一層的表示方式，控制流和數(shù)據(jù)流合一，更適合JIT優(yōu)化。
• ANF形式：和Thorin類似，都采用Graph IR，都消除了Scope。但是沒有采用Thorin IR的CPS形式，而是表達(dá)能力類似，更直觀也更易檢查的ANF形式。

總結(jié)，MindIR希望通過Functional的方式更方便的實(shí)現(xiàn)自動(dòng)微分和隱式并行分析，Graph Based方式把控制流和數(shù)據(jù)流合一支持更高效的JIT優(yōu)化。

1、MindIR的詳解

參考[7]

MindIR文法繼承于ANF，其定義如下所示：

 ::=  |  ::= Parameter
 ::= Scalar | Named | Tensor | Type | Shape
               | Primitive | MetaFuncGraph | FuncGraph 
 ::= ( …)
 ::=  | 

MindIR中的ANode對(duì)應(yīng)于ANF的原子表達(dá)式，ANode有兩個(gè)子類分別為ValueNode和ParameterNode。ValueNode表示常數(shù)節(jié)點(diǎn)，可承載一個(gè)常數(shù)值（標(biāo)量、符號(hào)、張量、類型、維度等），也可以是一個(gè)原語函數(shù)（Primitive）或一個(gè)元函數(shù)（MetaFuncGraph）或一個(gè)普通函數(shù)（FuncGraph），因?yàn)樵诤瘮?shù)式編程中函數(shù)定義本身也是一個(gè)值。ParameterNode是參數(shù)節(jié)點(diǎn)，表示函數(shù)的形參。

MindIR中CNode對(duì)應(yīng)于ANF的復(fù)合表達(dá)式，表示一次函數(shù)調(diào)用。

在MindSpore自動(dòng)微分時(shí)，會(huì)計(jì)算ParameterNode和CNode的梯度貢獻(xiàn)，并返回最終ParameterNode的梯度，而不計(jì)算ValueNode的梯度。

下面以一段程序作為示例，對(duì)比理解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代碼對(duì)應(yīng)的ANF表達(dá)為：

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

對(duì)應(yīng)的MindIR為ir.dot：

在MindIR中，一個(gè)函數(shù)圖（FuncGraph）表示一個(gè)普通函數(shù)的定義，函數(shù)圖一般由ParameterNode、ValueNode和CNode組成有向無環(huán)圖，可以清晰地表達(dá)出從參數(shù)到返回值的計(jì)算過程。在上圖中可以看出，python代碼中兩個(gè)函數(shù)test/_f和func轉(zhuǎn)換成了兩個(gè)函數(shù)圖，其參數(shù)x和y轉(zhuǎn)換為函數(shù)圖的ParameterNode，每一個(gè)表達(dá)式轉(zhuǎn)換為一個(gè)CNode。CNode的第一個(gè)輸入鏈接著調(diào)用的函數(shù)，例如圖中的add、func、return。值得注意的是這些節(jié)點(diǎn)均是ValueNode，因?yàn)樗鼈儽焕斫鉃槌?shù)函數(shù)值。CNode的其他輸入鏈接這調(diào)用的參數(shù)，參數(shù)值可以來自于ParameterNode、ValueNode和其他CNode。

在ANF中每個(gè)表達(dá)式都用let表達(dá)式綁定為一個(gè)變量，通過對(duì)變量的引用來表示對(duì)表達(dá)式輸出的依賴，而在MindIR中每個(gè)表達(dá)式都綁定為一個(gè)節(jié)點(diǎn)，通過節(jié)點(diǎn)與節(jié)點(diǎn)之間的有向邊表示依賴關(guān)系。

2、函數(shù)式語義

MindIR較傳統(tǒng)計(jì)算圖的一個(gè)重要特性是不僅可以表達(dá)算子之間的數(shù)據(jù)依賴，還可以表達(dá)豐富的函數(shù)式語義。

高階函數(shù)

在MindIR中，函數(shù)的定義是由一個(gè)子圖來定義，但其本身可以是一個(gè)被傳遞的值，作為其他高階函數(shù)的輸入或輸出。 例如下面一個(gè)簡(jiǎn)單的示例中，函數(shù)f作為參數(shù)傳入了函數(shù)g，因此函數(shù)g是一個(gè)接收函數(shù)輸入的高階函數(shù)，函數(shù)f真正的調(diào)用點(diǎn)是在函數(shù)g內(nèi)部。

@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

對(duì)應(yīng)的MindIR為hof.dot：

在實(shí)際網(wǎng)絡(luò)訓(xùn)練腳本中，自動(dòng)求導(dǎo)泛函GradOperation和優(yōu)化器中常用到的Partial和HyperMap都是典型的高階函數(shù)。高階語義極大地提升了MindSpore表達(dá)的靈活性和簡(jiǎn)潔性。

控制流

控制流在MindIR中是以高階函數(shù)選擇調(diào)用的形式表達(dá)。這樣的形式把控制流轉(zhuǎn)換為高階函數(shù)的數(shù)據(jù)流，從而使得自動(dòng)微分算法更加強(qiáng)大。不僅可以支持?jǐn)?shù)據(jù)流的自動(dòng)微分，還可以支持條件跳轉(zhuǎn)、循環(huán)和遞歸等控制流的自動(dòng)微分。

下面以一個(gè)簡(jiǎn)單的斐波那契用例來演示說明。

@ms_function
def fibonacci(n):
 if(n < 1):
 return 0
 elif(n == 1):
 return 1
 else:
 return fibonacci(n-1) + fibonacci(n-2)

對(duì)應(yīng)的MindIR為cf.dot：

其中fibonacci是頂層函數(shù)圖，在頂層中有兩個(gè)函數(shù)圖被switch選擇調(diào)用。?fibonacci是第一個(gè)if的True分支，?fibonacci是第一個(gè)if的False分支。在?fibonacci中被調(diào)用的??fibonacci是elif的True分支，??fibonacci是elif的False分支。這里需要理解的關(guān)鍵是在MindIR中，條件跳轉(zhuǎn)和遞歸是以高階控制流的形式表達(dá)的。例如，?fibonacci和?fibonacci是作為switch算子的參數(shù)傳入，switch根據(jù)條件參數(shù)選擇哪一個(gè)函數(shù)作為返回值。因此，switch是把輸入的函數(shù)當(dāng)成普通的值做了一個(gè)二元選擇操作，并沒有調(diào)用，而真正的函數(shù)調(diào)用是在緊隨switch后的CNode上完成。

自由變量和閉包

閉包（closure）是一種編程語言特性，它指的是代碼塊和作用域環(huán)境的結(jié)合。自由變量(free variable)是指在代碼塊中引用作用域環(huán)境中的變量而非局部變量。在MindIR中，代碼塊是以函數(shù)圖呈現(xiàn)的，而作用域環(huán)境可以理解為該函數(shù)被調(diào)用時(shí)的上下文環(huán)境，自由變量的捕獲方式是值拷貝而非引用。

一個(gè)典型的閉包用例如下：

@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

對(duì)應(yīng)的MindIR為closure.dot：

在例子中，a和b是自由變量，因?yàn)閒unc/_inner中變量a和b是引用的其父圖func/_outer中定義的參數(shù)。變量closure是一個(gè)閉包，它是函數(shù)func/_inner與其上下文func/_outer(1, 2)的結(jié)合。因此，out1的結(jié)果是4，因?yàn)槠涞葍r(jià)于1+2+1，out2的結(jié)果是5，因?yàn)槠涞葍r(jià)于1+2+2。

審核編輯 黃昊宇