TVM設(shè)計與構(gòu)架構(gòu)建
本文檔適用于希望了解TVM體系結(jié)構(gòu)和/或在項目上進行積極開發(fā)的開發(fā)人員。該頁面的組織如下：
? 實例編譯流程Example Compilation Flow，描述TVM把一個模型的高級描述到可部署模塊的步驟。
? “邏輯體系結(jié)構(gòu)組件” Logical Architecture Components部分描述了邏輯組件。針對每個邏輯組件的特定內(nèi)容，按組件名稱組織。
? 開發(fā)人員操作手冊，以獲取有用的開發(fā)技巧。
本文提供了一些體系結(jié)構(gòu)的補充視圖。首先，回顧一個端到端的編譯流程，討論關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)和轉(zhuǎn)換。這個基于運行時runtime的視圖，著重于運行編譯器時，每個組件的交互。回顧代碼庫的邏輯模塊及其關(guān)系。這部分提供了設(shè)計的靜態(tài)總體視圖。
編譯流程示例
本節(jié)將研究編譯器中的示例編譯流程。下圖顯示了流程。從高層次上，包含幾個步驟：
? 導(dǎo)入：前端組件將模型提取到IRModule中，該模塊包含內(nèi)部表示模型的函數(shù)的集合。
? 轉(zhuǎn)換：編譯器將IRModule，轉(zhuǎn)換為另一個功能上等效，或近似等效的（例如，在量化的情況下）IRModule。許多轉(zhuǎn)換都是獨立于目標(biāo)（后端）的。允許目標(biāo)影響轉(zhuǎn)換管道的配置。
? 目標(biāo)翻譯：編譯器將IRModule轉(zhuǎn)換（代碼生成）為目標(biāo)指定的可執(zhí)行格式。目標(biāo)翻譯結(jié)果，封裝為runtime.Module，可以導(dǎo)出，加載和執(zhí)行，對目標(biāo)運行時runtime環(huán)境。
? 運行時runtime執(zhí)行：用戶load了runtime.Module，在支持的運行環(huán)境和運行編譯功能。

關(guān)鍵數(shù)據(jù)結(jié)構(gòu)
設(shè)計和理解復(fù)雜系統(tǒng)的最佳方法之一，識別關(guān)鍵數(shù)據(jù)結(jié)構(gòu)和算子（轉(zhuǎn)換）。這些數(shù)據(jù)結(jié)構(gòu)的API，一旦確定了關(guān)鍵數(shù)據(jù)結(jié)構(gòu)，便可以將系統(tǒng)分解為邏輯組件，這些邏輯組件可以定義關(guān)鍵數(shù)據(jù)結(jié)構(gòu)的集合，或數(shù)據(jù)結(jié)構(gòu)之間的轉(zhuǎn)換。
IRModule是整個堆棧中使用的主要數(shù)據(jù)結(jié)構(gòu)。IRModule（中間表示模塊）包含函數(shù)的集合。支持兩種主要的功能變體。
? relay :: Function是高級功能代碼表示。relay.Function通常對應(yīng)于端到端模型。可以將relay作為計算圖，額外支持控制流，遞歸和復(fù)雜的數(shù)據(jù)結(jié)構(gòu)。
? tir :: PrimFunc是一種低級代碼表示，包含一些元素，包括循環(huán)嵌套選擇，多維加載/存儲，線程和向量/張量指令。通常用于表示執(zhí)行模型中（可能是融合的）層的算子代碼。
在編譯期間，可以將relay函數(shù)降低為多個tir :: PrimFunc函數(shù)和一個調(diào)用這些tir :: PrimFunc函數(shù)的頂級函數(shù)。
轉(zhuǎn)換
已經(jīng)涵蓋了關(guān)鍵數(shù)據(jù)結(jié)構(gòu)，談?wù)勣D(zhuǎn)換。每個轉(zhuǎn)換可以滿足以下目的之一：
? 優(yōu)化：將代碼轉(zhuǎn)換為等效的，可能更優(yōu)化的版本。
? 降維：將代碼轉(zhuǎn)換為更接近目標(biāo)的，較低層表示。
Relay/轉(zhuǎn)換包含優(yōu)化模型的通道集合。這些優(yōu)化包括常見的代碼優(yōu)化（如常量折疊和死代碼消除），以及張量計算特定的遍歷（如布局轉(zhuǎn)換和縮放因子折疊）。
在Relay優(yōu)化管道的末端附近，運行pass（FuseOps），將端到端功能（例如MobileNet），劃分為子功能（例如conv2d-relu）段，稱這些功能為段。可將原始問題分為兩個子問題：
? 每個子功能的編譯和優(yōu)化。
? 總體執(zhí)行結(jié)構(gòu)：需要對所生成的子函數(shù)，進行一系列調(diào)用，執(zhí)行整個模型。
使用低級的Tir，編譯和優(yōu)化每個子功能。對于特定目標(biāo)，可以直接進入目標(biāo)平移，使用外部代碼生成器。
有幾種不同的方法（在Relay/后端），處理對整個執(zhí)行問題的調(diào)用。對于具有已知形狀，無控制流的簡單模型，可以降低到將執(zhí)行結(jié)構(gòu)，存儲在圖中的圖運行時runtime。支持虛擬機后端進行動態(tài)執(zhí)行。支持提前編譯，該編譯將高級執(zhí)行結(jié)構(gòu)，編譯為可執(zhí)行文件和生成的原始函數(shù)。所有這些執(zhí)行模式，都由統(tǒng)一的runtime.Module 接口封裝 ，將在本文的后面部分中進行討論。
tir / transform包含用于TIR級別功能的轉(zhuǎn)換過程。許多Tir通道的目的是降維。例如，可以通過多種途徑，將多維接入到一維指針訪問，將內(nèi)在函數(shù)擴展為特定于目標(biāo)的內(nèi)在函數(shù)，修飾函數(shù)條目，滿足運行時runtime調(diào)用約定。有一些優(yōu)化過程，如簡化訪問索引和消除無效代碼。
LLVM，CUDA C和其它目標(biāo)編譯器，可以在目標(biāo)處理許多低級優(yōu)化。將低級優(yōu)化（例如寄存器分配），留給了下游編譯器，只專注于未涵蓋的優(yōu)化。
搜索空間和基于學(xué)習(xí)的轉(zhuǎn)換
描述的轉(zhuǎn)換過程是確定性的，基于規(guī)則的。TVM堆棧的一個設(shè)計目標(biāo)，支持針對不同硬件平臺的高性能代碼優(yōu)化。將需要研究盡可能多的優(yōu)化選擇，包括但不限于多維張量訪問，循環(huán)切片行為，特殊的加速器內(nèi)存層次結(jié)構(gòu)和線程化。
很難定義做出所有選擇的試探法。將采用基于搜索和學(xué)習(xí)的方法。首先定義可以用來轉(zhuǎn)換代碼操作的集合。示例操作包括循環(huán)轉(zhuǎn)換，內(nèi)聯(lián)，向量化。稱這些操作為調(diào)度原語。調(diào)度原語的集合定義了，可以對代碼進行的可能優(yōu)化的搜索空間。系統(tǒng)搜索不同的可能調(diào)度序列，以選擇最佳調(diào)度組合。搜索過程通常以機器學(xué)習(xí)算法為原則。
搜索完成后，可以為（可能是融合的）算子，記錄最佳調(diào)度順序。編譯器可以僅查找最佳調(diào)度序列，應(yīng)用于代碼。值得注意的是，此調(diào)度應(yīng)用代碼布局，完全類似于基于規(guī)則的變換，能夠與傳統(tǒng)流程，共享相同的接口約定。
使用基于搜索的優(yōu)化，處理最初的Tir函數(shù)生成問題。此模塊部分稱為AutoTVM（auto_scheduler）。隨著繼續(xù)開發(fā)TVM堆棧，希望將基于學(xué)習(xí)的轉(zhuǎn)換擴展到更多領(lǐng)域。
目標(biāo)轉(zhuǎn)換
目標(biāo)轉(zhuǎn)換階段，將IRModule轉(zhuǎn)換為相應(yīng)的目標(biāo)可執(zhí)行格式。對于x86和ARM等后端，使用LLVM IRBuilder，構(gòu)建內(nèi)存中的LLVM IR。可以生成諸如CUDA C和OpenCL之類的，源代碼級語言。支持通過外部代碼生成器，將Relay函數(shù)（子圖）直接轉(zhuǎn)換為特定目標(biāo)。最終代碼生成階段應(yīng)盡可能輕巧。絕大部分的轉(zhuǎn)換和降維，都應(yīng)在目標(biāo)轉(zhuǎn)換之前進行。
提供了一個Target結(jié)構(gòu)，指定編譯目標(biāo)。目標(biāo)翻譯階段之前的轉(zhuǎn)換，可能受到目標(biāo)的影響-例如，目標(biāo)的向量長度，改變向量化行為。
運行時runtime執(zhí)行
TVM運行時runtime的主要目標(biāo)，提供一個最小的API，使用選擇的語言（包括Python，C ++，Rust，Go，Java和JavaScript），加載和執(zhí)行已編譯的工件。下面的代碼片段，顯示了Python中示例：
import tvm
# Example runtime execution program in python, with type annotated
mod: tvm.runtime.Module = tvm.runtime.load_module(“compiled_artifact.so”)
arr: tvm.runtime.NDArray = tvm.nd.array([1, 2, 3], ctx=tvm.gpu(0))
fun: tvm.runtime.PackedFunc = mod[“addone”]
fun(a)
print(a.asnumpy())
tvm.runtime.Module封裝編譯結(jié)果。runtime.Module包含一個GetFunction方法，用于按名稱獲取PackedFuncs。
tvm.runtime.PackedFunc，兩個生成的函數(shù)的類型清理的函數(shù)接口。runtime.PackedFunc可采用以下類型的參數(shù)并返回值：POD類型（int，float），字符串，runtime.PackedFunc，runtime.Module，runtime.NDArray以及runtime.Object的其它子類。
tvm.runtime.Module，tvm.runtime.PackedFunc，將運行時runtime模塊化的強大機制。在CUDA上獲得上述addone函數(shù)，可以使用LLVM生成主機端代碼，計算啟動參數(shù)（例如線程組的大小），從CUDAModule，調(diào)用另一個由PackedFunc支持的PackedFunc。 CUDA驅(qū)動代碼API。相同的機制可用于OpenCL內(nèi)核。
上面的示例僅處理簡單的addone函數(shù)。下面的代碼段給出了使用同一接口執(zhí)行端到端模型的示例：
import tvm
# Example runtime execution program in python, with types annotated
factory: tvm.runtime.Module = tvm.runtime.load_module(“resnet18.so”)
# Create a stateful graph execution module for resnet18 on gpu(0)
gmod: tvm.runtime.Module = factory"resnet18"
data: tvm.runtime.NDArray = get_input_data()
# set input
gmod[“set_input”](0, data)
# execute the model
gmod"run"
# get the output
result = gmod"get_output".asnumpy()
主要優(yōu)點，runtime.Module和runtime.PackedFunc足以封裝算子級代碼（例如addone），及端到端模型。
總結(jié)與討論
總之，編譯流程中的關(guān)鍵數(shù)據(jù)結(jié)構(gòu)為：
? IRModule：包含relay.Function和tir.PrimFunc
? runtime.Module：包含runtime.PackedFunc
編譯的大部分內(nèi)容是關(guān)鍵數(shù)據(jù)結(jié)構(gòu)之間的轉(zhuǎn)換。
? relay / transform和tir / transform，基于規(guī)則的確定性轉(zhuǎn)換
? auto_scheduler和autotvm，包含基于搜索的轉(zhuǎn)換
編譯流程示例，只是TVM堆棧的典型用例。將這些關(guān)鍵數(shù)據(jù)結(jié)構(gòu)和轉(zhuǎn)換，開放給python和C ++ API。除了感興趣的數(shù)據(jù)結(jié)構(gòu)，從numpy.ndarray更改為tvm.IRModule之外，可以像使用numpy，一樣使用TVM。以下是一些用例示例：
? 使用python API直接構(gòu)造IRModule。
? 組成一組自定義的轉(zhuǎn)換（例如，自定義量化）。
? 使用TVM的python API直接操作IR。
邏輯架構(gòu)組件

TVM體系結(jié)構(gòu)圖
上圖顯示了項目中的主要邏輯組件。獲取有關(guān)組件及關(guān)系的信息。
tvm/support
支持模塊包含最常用的基礎(chǔ)構(gòu)建實用代碼，例如通用競技場分配器，套接字和日志記錄。
tvm/runtime
運行時runtime是TVM堆棧的基礎(chǔ)。提供了加載和執(zhí)行已編譯工件的機制。運行時runtime定義了一組穩(wěn)定的標(biāo)準(zhǔn)C API，以與諸如Python和Rust的前端語言進行接口。
除了runtime :: PackedFunc之外，runtime :: Object是TVM運行時runtime中的，主要數(shù)據(jù)結(jié)構(gòu)之一。帶有類型索引的引用計數(shù)基類，支持運行時runtime類型檢查和向下轉(zhuǎn)換。目標(biāo)系統(tǒng)允許開發(fā)人員，向運行時runtime引入新的數(shù)據(jù)結(jié)構(gòu)，例如數(shù)組，映射和新的IR數(shù)據(jù)結(jié)構(gòu)。
除了部署用例之外，編譯器本身還大量使用TVM的運行時runtime機制。所有的IR數(shù)據(jù)結(jié)構(gòu)都是runtime :: Object的子類，可以從Python前端直接訪問和操作。使用PackedFunc機制，將各種API開放給前端。
在運行時runtime的子目錄（例如runtime / opencl）中，定義了對不同硬件后端的運行時runtime支持。這些特定于硬件的運行時runtime模塊定義，用于設(shè)備內(nèi)存分配和設(shè)備功能序列化的API。
runtime / rpc為PackedFunc實現(xiàn)RPC支持。可以使用RPC機制，將交叉編譯的庫，發(fā)送到遠程設(shè)備，確定執(zhí)行性能的基準(zhǔn)。rpc基礎(chǔ)架構(gòu)，支持從廣泛的硬件后端收集數(shù)據(jù)，進行基于學(xué)習(xí)的優(yōu)化。
? TVM運行系統(tǒng)
? 調(diào)試器
? 將VM放入TVM：Relay虛擬機
? 模塊序列化簡介
tvm/node
節(jié)點模塊在runtime :: Object的基礎(chǔ)上，為IR數(shù)據(jù)結(jié)構(gòu)添加了其它功能。主要包括反射，序列化，結(jié)構(gòu)等效和散列。
使用了節(jié)點模塊，可以通過在Python中的名稱，直接訪問TVM的IRNode的任何字段。
x = tvm.tir.Var(“x”, “int32”)
y = tvm.tir.Add(x, x)
# a and b are fields of a tir.Add node
# we can directly use the field name to access the IR structures
assert y.a == x
可以將任意IR節(jié)點，序列化為JSON格式，然后將其加載回。保存/存儲和檢查IR節(jié)點的能力，使編譯器更易于訪問，提供了基礎(chǔ)。
tvm/ir
在TVM / IR文件夾，包含跨所有IR功能變體的，統(tǒng)一的數(shù)據(jù)結(jié)構(gòu)和接口。tvm / ir中的組件由tvm / relay和tvm / tir共享，包括
? ir模塊
? 類型
? PassContext和Pass
? 算子
功能的不同變體（例如relay.Function和tir.PrimFunc），可以共存于IRModule中。盡管這些變體可能不具有相同的內(nèi)容表示，使用相同的數(shù)據(jù)結(jié)構(gòu)來表示類型。使用相同的數(shù)據(jù)結(jié)構(gòu)，表示這些變量的功能（類型）名稱。一旦明確定義了調(diào)用約定，統(tǒng)一類型系統(tǒng)，允許一個函數(shù)變換，調(diào)用另一個函數(shù)。這為將來的跨功能變量優(yōu)化，打開了大門。
提供了一個統(tǒng)一的PassContext，用于配置傳遞算子，提供了通用的復(fù)合pass，執(zhí)行pass管道。以下代碼段給出了PassContext配置的示例。
# configure the behavior of the tir.UnrollLoop pass
with tvm.transform.PassContext(config={“tir.UnrollLoop”: { “auto_max_step”: 10 }}):
# code affected by the pass context
Op是表示所有系統(tǒng)定義的原始算子/內(nèi)部函數(shù)的通用類。開發(fā)人員可以向系統(tǒng)注冊新的Op，以及其它屬性（例如Op是否是元素化的）。
? Pass Infrastructure
tvm/target
目標(biāo)模塊包含將IRModule轉(zhuǎn)換為目標(biāo)runtime.Module的，所有代碼生成器。提供了描述目標(biāo)的通用Target類。
通過查詢目標(biāo)中的屬性信息和注冊到每個目標(biāo)id（cuda，opencl）的內(nèi)置信息，可以根據(jù)目標(biāo)定制編譯管道。
tvm/ir
TIR包含低級代碼表示的定義。使用tir :: PrimFunc表示，可以通過TIR傳遞轉(zhuǎn)換的函數(shù)。除IR數(shù)據(jù)結(jié)構(gòu)外，tir模塊還通過公共Op注冊表，以及tir / transform中的轉(zhuǎn)換，傳遞定義了一組內(nèi)置的內(nèi)在函數(shù)及其屬性。
tvm/arith
此模塊與TIR緊密相關(guān)。低級代碼生成中的關(guān)鍵問題之一，分析索引的算術(shù)屬性-正則性，變量邊界以及描述迭代空間的整數(shù)集。arith模塊提供了一組進行（主要是整數(shù)）分析的工具。TIR pass可以使用這些分析，簡化和優(yōu)化代碼。
tvm/te
te代表“張量表達式”。這是一個特定領(lǐng)域的語言模塊，允許通過編寫張量表達式，快速構(gòu)建tir :: PrimFunc變體。重要的是，張量表達式本身不是可以存儲到IRModule中的自包含函數(shù)。相反，是IR的一個片段，拼接在一起，構(gòu)建IRModule。
te / schedule提供了一組調(diào)度原語，控制所生成的功能。可能會將一些調(diào)度組件引入tir :: PrimFunc本身。
? InferBound Pass
? 混合前端開發(fā)人員指南
? InferBound Pass
? Hybrid Frontend Developer Guide
tvm/topi
可以針對每個用例直接通過TIR，或張量表達式（TE），構(gòu)造算子。 topi（張量算子清單）提供了一組預(yù)定義的算子（在TE或TIR中），由numpy定義，在常見的深度學(xué)習(xí)工作負(fù)載中找到。提供了一組公共調(diào)度模板，在不同目標(biāo)平臺上，獲得高性能的實現(xiàn)。
tvm/relay
Relay是用于表示完整模型的高級功能性IR。在relay.transform中定義了各種優(yōu)化。Relay編譯器定義了多種語言，每種語言旨在支持特定的優(yōu)化樣式。值得注意的是QNN（用于導(dǎo)入預(yù)量化模型），VM（用于降級為動態(tài)虛擬機），內(nèi)存（用于內(nèi)存優(yōu)化）。
? Relay IR簡介
? Relay算子策略
? 轉(zhuǎn)換布局通道
? Introduction to Relay IR
? Relay Operator Strategy
? Convert Layout Pass
tvm/autotvm
AutoTVM和AutoScheduler，都是自動進行基于搜索的代碼優(yōu)化的組件。主要包括：
? Cost models和特征提取。
? 一種記錄格式，用于存儲調(diào)度基準(zhǔn)結(jié)果，進行cost model構(gòu)建。
? 一組有關(guān)代碼轉(zhuǎn)換的搜索策略。
? Cost models and feature extraction.
? A record format for storing program benchmark results for cost model construction.
? A set of search policies over program transformations.
自動化代碼優(yōu)化，仍然是活躍的研究領(lǐng)域。試圖對設(shè)計進行模塊化，研究人員可以通過Python綁定，快速修改組件或自己的應(yīng)用算法，自定義搜索，從Python綁定中，插入算法。
? 基準(zhǔn)性能日志格式Benchmark Performance Log Format
前端
前端將來自不同框架的模型，存放在TVM堆棧中。 tvm.relay.frontend是模型提取API的命名空間。
? TensorFlow前端

總結(jié)

以上是生活随笔為你收集整理的TVM设计与构架构建的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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