Relay IR表示

Relay IR簡(jiǎn)介
本節(jié)介紹了 Relay IR——第二代 NNVM。期待兩種背景的讀者——具有編程語(yǔ)言背景的讀者和熟悉計(jì)算圖表示的，深度學(xué)習(xí)框架開發(fā)人員。
簡(jiǎn)要總結(jié)了設(shè)計(jì)目標(biāo)，討論要點(diǎn)。
? 支持傳統(tǒng)的數(shù)據(jù)流式編程和轉(zhuǎn)換。
? 支持函數(shù)式作用域，let-binding，成為功能齊全的可微分語(yǔ)言。
? 能夠允許用戶混合兩種編程風(fēng)格。
使用 Relay 構(gòu)建計(jì)算圖
傳統(tǒng)的深度學(xué)習(xí)框架，使用計(jì)算圖作為中間表示。計(jì)算圖（或數(shù)據(jù)流圖）是表示計(jì)算的有向無(wú)環(huán)圖 (DAG)。由于缺乏控制流，數(shù)據(jù)流圖表達(dá)的計(jì)算受限制，容易實(shí)現(xiàn)自動(dòng)微分，異構(gòu)執(zhí)行環(huán)境編譯（例如，專用硬件上執(zhí)行圖）。
可以使用 Relay 構(gòu)建計(jì)算（數(shù)據(jù)流）圖。如何構(gòu)造一個(gè)簡(jiǎn)單的二節(jié)點(diǎn)圖。示例語(yǔ)法與 NNVMv1 等現(xiàn)有計(jì)算圖 IR，唯一區(qū)別：
? 現(xiàn)有框架通常使用圖和子圖
? Relay 使用函數(shù) eg – 表示圖形fn (%x)
每個(gè)數(shù)據(jù)流節(jié)點(diǎn)，都是 Relay 中的一個(gè) CallNode。Relay Python DSL快速構(gòu)建數(shù)據(jù)流圖。構(gòu)造Add節(jié)點(diǎn)，兩個(gè)輸入點(diǎn)都指向%1。深度學(xué)習(xí)框架評(píng)估時(shí)，按照拓?fù)漤樞蛴?jì)算節(jié)點(diǎn)， %1只會(huì)計(jì)算一次。實(shí)現(xiàn)訪問者，打印結(jié)果，嵌套的 Call，變成log(%x) + log(%x)。
DAG 中存在共享節(jié)點(diǎn)時(shí)，對(duì)語(yǔ)義的不同解釋造成的。在普通的函數(shù)式編程 IR 中，嵌套表達(dá)式視為表達(dá)式樹，不考慮%1在%2中，實(shí)際重用了兩次。
Relay IR 到了這種差異。通常，深度學(xué)習(xí)框架用戶，用這種方式構(gòu)建計(jì)算圖，經(jīng)常發(fā)生 DAG 節(jié)點(diǎn)重用。用文本格式打印 Relay時(shí)，每行打印一個(gè) CallNode，將每個(gè) CallNode分配一個(gè)臨時(shí) id (%1, %2)，可以引用每個(gè)公共節(jié)點(diǎn)。
模塊：支持多種功能（圖形）
介紹了如何將數(shù)據(jù)流圖構(gòu)建為函數(shù)。能不能支持多種功能，可以互相調(diào)用？Relay將多個(gè)功能組合在一個(gè)模塊中；下面的代碼顯示，一個(gè)函數(shù)調(diào)用另一個(gè)函數(shù)的示例。
def @muladd(%x, %y, %z) {
%1 = mul(%x, %y)
%2 = add(%1, %z)
%2
}
def @myfunc(%x) {
%1 = @muladd(%x, 1, 2)
%2 = @muladd(%1, 2, 3)
%2
}
該模塊可以被視為一個(gè)Map<GlobalVar, Function>，這里 GlobalVar 只是一個(gè) id，表示模塊中的功能。@muladd與@myfunc是上面例子中的 GlobalVars。當(dāng)CallNode調(diào)用另一個(gè)函數(shù)時(shí)，對(duì)應(yīng)的 GlobalVar 存儲(chǔ)在 CallNode 的 op 字段中。包含一個(gè)間接級(jí)別——使用相應(yīng)的 GlobalVar，查找調(diào)用函數(shù)的主體。可以直接將函數(shù)的引用，存儲(chǔ)為 CallNode 中的 op。為什么需要引入 GlobalVar 呢？ GlobalVar 解耦了定義/聲明，啟用了函數(shù)的遞歸和延遲聲明。
def @myfunc(%x) {
%1 = equal(%x, 1)
if (%1) {
%x
} else {
%2 = sub(%x, 1)
%3 = @myfunc(%2)
%4 = add(%3, %3)
%4
}
}
在上面的例子中，@myfunc遞歸地調(diào)用。使用 GlobalVar@myfunc，表示函數(shù)避免了數(shù)據(jù)結(jié)構(gòu)中的循環(huán)依賴。Relay 相對(duì)于 NNVMv1，有以下改進(jìn)：
? 簡(jiǎn)潔的文本格式，可簡(jiǎn)化寫入passes的調(diào)試。
? 在聯(lián)合模塊中，對(duì)子圖函數(shù)的一流支持，使得聯(lián)合優(yōu)化會(huì)成為可能，如內(nèi)聯(lián)和調(diào)用約定規(guī)范。
? 簡(jiǎn)單的前端語(yǔ)言互操作，所有數(shù)據(jù)結(jié)構(gòu)，可以在 Python 中訪問，允許在 Python 中，快速構(gòu)建優(yōu)化原型，與 C++ 代碼混合。
讓綁定和作用域
深度學(xué)習(xí)框架中，使用的舊方法構(gòu)建計(jì)算圖。將討論 Relay 引入新結(jié)構(gòu)——let 綁定。
每一種高級(jí)編程語(yǔ)言，使用 let 綁定。在 Relay 中，Let(var, value, body)表示三個(gè)字段的數(shù)據(jù)結(jié)構(gòu)。評(píng)估let表達(dá)式，先評(píng)估 value 部分，分配給 var，在 body 表達(dá)式中，返回評(píng)估結(jié)果。
可以使用一系列 let 綁定，構(gòu)造邏輯上等效數(shù)據(jù)流程序。
嵌套的 let 綁定稱為 A 范式，通常用作函數(shù)式編程語(yǔ)言中的 IR，仔細(xì)看看 AST 結(jié)構(gòu)。雖然這兩個(gè)程序在語(yǔ)義上相同（除了 A 范式，有 let 前綴），但AST 結(jié)構(gòu)不同。
程序優(yōu)化采用這些 AST 數(shù)據(jù)結(jié)構(gòu)，進(jìn)行轉(zhuǎn)換，兩種不同的結(jié)構(gòu)，將影響要編寫的編譯器代碼。例如，如果檢測(cè)一個(gè)模式：add(log(x),y)。
? 在data-flow形式中，先訪問add節(jié)點(diǎn)，直接查看第一個(gè)參數(shù)，是不是日志
? 在 A 范式中，不能直接檢查，要添加的第一個(gè)輸入是%v1–需要保留從變量到綁定值的映射查找， %v1是一個(gè)日志。
不同的數(shù)據(jù)結(jié)構(gòu)，影響編寫轉(zhuǎn)換的方式。為什么需要 let 綁定？PL 是一個(gè)相當(dāng)成熟的領(lǐng)域。
為什么可能需要綁定
let 綁定的一個(gè)關(guān)鍵用法，指定計(jì)算范圍。下面的例子，沒有使用 let 綁定。
試圖決定應(yīng)該在哪里評(píng)估 node 時(shí)，問題就出現(xiàn)了%1。特別是，雖然文本格式似乎建議應(yīng)該%1，在if 范圍之外評(píng)估節(jié)點(diǎn)，但 AST（如圖所示）不建議這樣做。實(shí)際上，數(shù)據(jù)流圖從未定義評(píng)估范圍。這在語(yǔ)義中，引入了一些歧義。
有閉包時(shí)，這種歧義變得更加有趣。考慮以下程序，返回一個(gè)閉包。不知道應(yīng)該在哪里計(jì)算%1；可以在閉包內(nèi)部或外部。
fn (%x) {
%1 = log(%x)
%2 = fn(%y) {
add(%y, %1)
}
%2
}
let 綁定解決了這個(gè)問題，值的計(jì)算發(fā)生在 let 節(jié)點(diǎn)。在這兩個(gè)程序中，如果%1 = log(%x)更改成let %v1 = log(%x)，明確指定計(jì)算位置，在 if 范圍和閉包外。let-binding 提供了更精確的計(jì)算站點(diǎn)規(guī)范，生成后端代碼時(shí)，可能很有用（此類規(guī)范在 IR 中）。
另一方面，不指定計(jì)算范圍的數(shù)據(jù)流形式，確實(shí)有優(yōu)勢(shì)——在生成代碼時(shí)，無(wú)需擔(dān)心將 let 放在哪里。數(shù)據(jù)流表單為后面的passes，決定將評(píng)估點(diǎn)放在哪里。在優(yōu)化的初始階段，當(dāng)發(fā)現(xiàn)方便時(shí)，使用程序的數(shù)據(jù)流形式，可能不是一個(gè)壞主意。Relay 中的許多優(yōu)化，都是為了優(yōu)化數(shù)據(jù)流程序編寫的。
將 IR 降低到實(shí)際運(yùn)行時(shí)程序時(shí)，需要精確計(jì)算的范圍。使用子函數(shù)和閉包時(shí)，希望明確指定計(jì)算范圍，應(yīng)該發(fā)生在哪里。Let-binding 可用于后期執(zhí)行特定優(yōu)化，解決問題。
對(duì) IR 轉(zhuǎn)換的影響
大多數(shù)函數(shù)式編程語(yǔ)言，用 A 范式分析，不需要表達(dá)式DAG。
Relay同時(shí)支持?jǐn)?shù)據(jù)流形式和 let 綁定。讓框架開發(fā)人員熟悉的表示。如何編寫 pass，有一些影響：
? 如果來(lái)自數(shù)據(jù)流背景，要處理 let，將 var 映射到表達(dá)式，遇到 var 時(shí)，執(zhí)行查找。需要一個(gè)從表達(dá)式到轉(zhuǎn)換表達(dá)式的映射。有效地刪除程序中的所有l(wèi)ets。
? 如果來(lái)自 PL 背景，喜歡 A 范式，將提供 A 范式pass的數(shù)據(jù)流。
? 對(duì)于 PL 人員，當(dāng)實(shí)現(xiàn)某些東西（例如，數(shù)據(jù)流到 ANF 的轉(zhuǎn)換）時(shí)，表達(dá)式可以是 DAG，應(yīng)該使用 Map<Expr, Result> 訪問表達(dá)式，計(jì)算一次轉(zhuǎn)換后的結(jié)果，結(jié)果表達(dá)式保持公共結(jié)構(gòu)。
有一些額外的高級(jí)概念，例如符號(hào)形狀推理，未涵蓋的多態(tài)函數(shù)。

參考鏈接：
https://tvm.apache.org/docs/dev/relay_intro.html

總結(jié)

以上是生活随笔為你收集整理的Relay IR表示的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。