關(guān)于AI框架中圖層IR的分析

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

AI框架圖層IR的分析

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

背景

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

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

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

業(yè)界IR的介紹

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

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

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

push 3
push a
multiply
push a
substract

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

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

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

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

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

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

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

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

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

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

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

def foo(x):
return x+1

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

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

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

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

比如如下的一段python代碼，求小于n的所有素數(shù)的積。

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

當(dāng)采用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)

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

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

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

參考【3】

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

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

它的文法定義如下，參考[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形式表達，如果翻譯為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形式簡單易懂。

AI框架中圖層IR的作用

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

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

• Eager執(zhí)行模式一般是利用宿主語言（現(xiàn)在主要是Python）的特性進行解釋執(zhí)行，里面使用了重載和Tape的一些技巧。
• Graph執(zhí)行模式主要是拿到AI模型的圖結(jié)構(gòu)，然后進行編譯優(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òu)圖，好處是與Eagle模式比較容易匹配，實現(xiàn)簡單，缺點是控制流的轉(zhuǎn)換比較麻煩、執(zhí)行序列如果與算子執(zhí)行結(jié)果相關(guān)的話不好實現(xiàn)、不容易處理副作用，所以TF的AutoGraph還需要結(jié)合AST分析解決控制流轉(zhuǎn)換的問題；第三種是AST JIT（Pytorch的TorchScript），基于Python的AST進行構(gòu)圖，優(yōu)點是轉(zhuǎn)換的功能可以比較全面，包括控制流等，缺點是實現(xiàn)復(fù)雜，許多Python動態(tài)特性實現(xiàn)起來工作量大。
• Staging執(zhí)行模式類似在Eager模式中，通過Python修飾符，對部分子圖進行編譯執(zhí)行加速（使用Tracing JIT或者AST JIT），也會用到圖IR。

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

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

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

• 張量表達

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

• 自動微分

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

• JIT能力

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

• 隱式并行

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

• Loop優(yōu)化

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

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

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

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

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

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

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

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

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

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

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

對應(yīng)的 Primal CFG中，插入了 %1 phi節(jié)點作為啞phi節(jié)點來記錄控制流。然后在AD CFG中使用此 %1 來進行控制（%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的自動微分最終是通過迭代的方式來實現(xiàn)的，帶Scope的SSA形式需要解決邊界傳遞的問題對自動微分還是會帶來一些處理上的麻煩。

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

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

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

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

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

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

如何實現(xiàn)自動微分？

對于非控制流，計算AD的方法和上述的BB塊內(nèi)計算AD的方法相同。對于控制流，函數(shù)式IR采用了不同的處理方式，將迭代轉(zhuǎn)換為遞歸，并且通過switch函數(shù)來進行分支的選擇。例如上述相同的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() (此時return 5*5*5)

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

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

/# 對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的好處是對自動微分友好，比較適合做并行分析，不過挑戰(zhàn)在于函數(shù)IR的副作用消除以及函數(shù)式IR在執(zhí)行態(tài)的性能（含有遞歸對執(zhí)行不友好）。

MindSpore的設(shè)計思考

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

• Functional—更自然的自動微分實現(xiàn)方式和更方便的隱式并行分析能力：函數(shù)作為一等公民，支持高階函數(shù)，包括控制流也通過特殊的函數(shù)來實現(xiàn)，可以以統(tǒng)一的形式來實現(xiàn)微分；函數(shù)以無副作用的方式實現(xiàn)，與命令式語言相比，可簡化分析和實現(xiàn)更多的優(yōu)化；原生支持閉包，一方面可以方便的表達用戶源代碼中的閉包表示，另外也可以自然的支持自動微分算法中在反向函數(shù)中要訪問原始函數(shù)的中間結(jié)果的要求：反向函數(shù)訪問中間結(jié)果，并且作為一個閉包返回；使用基于數(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形式，而是表達能力類似，更直觀也更易檢查的ANF形式。

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

1、MindIR的詳解

參考[7]

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

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

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

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

在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代碼對應(yīng)的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

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

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

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

2、函數(shù)式語義

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

高階函數(shù)

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

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

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

控制流

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

下面以一個簡單的斐波那契用例來演示說明。

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

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

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

自由變量和閉包

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

一個典型的閉包用例如下：

@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

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

在例子中，a和b是自由變量，因為func/_inner中變量a和b是引用的其父圖func/_outer中定義的參數(shù)。變量closure是一個閉包，它是函數(shù)func/_inner與其上下文func/_outer(1, 2)的結(jié)合。因此，out1的結(jié)果是4，因為其等價于1+2+1，out2的結(jié)果是5，因為其等價于1+2+2。

審核編輯 黃昊宇