TVM設計與構架構建
本文檔適用於希望了解TVM體系結構和/或在項目上進行積極開發的開發人員。該頁面的組織如下：
• 實例編譯流程Example Compilation Flow，描述TVM把一個模型的高級描述到可部署模塊的步驟。
• “邏輯體系結構組件” Logical Architecture Components部分描述了邏輯組件。針對每個邏輯組件的特定內容，按組件名稱組織。
• 開發人員操作手冊，以獲取有用的開發技巧。
本文提供了一些體系結構的補充視圖。首先，回顧一個端到端的編譯流程，討論關鍵的數據結構和轉換。這個基於運行時runtime的視圖，着重於運行編譯器時，每個組件的交互。回顧代碼庫的邏輯模塊及其關系。這部分提供了設計的靜態總體視圖。
編譯流程示例
本節將研究編譯器中的示例編譯流程。下圖顯示了流程。從高層次上，包含幾個步驟：
• 導入：前端組件將模型提取到IRModule中，該模塊包含內部表示模型的函數的集合。
• 轉換：編譯器將IRModule，轉換為另一個功能上等效，或近似等效的（例如，在量化的情況下）IRModule。許多轉換都是獨立於目標（后端）的。允許目標影響轉換管道的配置。
• 目標翻譯：編譯器將IRModule轉換（代碼生成）為目標指定的可執行格式。目標翻譯結果，封裝為runtime.Module，可以導出，加載和執行，對目標運行時runtime環境。
• 運行時runtime執行：用戶load了runtime.Module，在支持的運行環境和運行編譯功能。
關鍵數據結構
設計和理解復雜系統的最佳方法之一，識別關鍵數據結構和算子（轉換）。這些數據結構的API，一旦確定了關鍵數據結構，便可以將系統分解為邏輯組件，這些邏輯組件可以定義關鍵數據結構的集合，或數據結構之間的轉換。
IRModule是整個堆棧中使用的主要數據結構。IRModule（中間表示模塊）包含函數的集合。支持兩種主要的功能變體。
• relay :: Function是高級功能代碼表示。relay.Function通常對應於端到端模型。可以將relay作為計算圖，額外支持控制流，遞歸和復雜的數據結構。
• tir :: PrimFunc是一種低級代碼表示，包含一些元素，包括循環嵌套選擇，多維加載/存儲，線程和向量/張量指令。通常用於表示執行模型中（可能是融合的）層的算子代碼。
在編譯期間，可以將relay函數降低為多個tir :: PrimFunc函數和一個調用這些tir :: PrimFunc函數的頂級函數。
轉換
已經涵蓋了關鍵數據結構，談談轉換。每個轉換可以滿足以下目的之一：
• 優化：將代碼轉換為等效的，可能更優化的版本。
• 降維：將代碼轉換為更接近目標的，較低層表示。
Relay/轉換包含優化模型的通道集合。這些優化包括常見的代碼優化（如常量折疊和死代碼消除），以及張量計算特定的遍歷（如布局轉換和縮放因子折疊）。
在Relay優化管道的末端附近，運行pass（FuseOps），將端到端功能（例如MobileNet），划分為子功能（例如conv2d-relu）段，稱這些功能為段。可將原始問題分為兩個子問題：
• 每個子功能的編譯和優化。
• 總體執行結構：需要對所生成的子函數，進行一系列調用，執行整個模型。
使用低級的Tir，編譯和優化每個子功能。對於特定目標，可以直接進入目標平移，使用外部代碼生成器。
有幾種不同的方法（在Relay/后端），處理對整個執行問題的調用。對於具有已知形狀，無控制流的簡單模型，可以降低到將執行結構，存儲在圖中的圖運行時runtime。支持虛擬機后端進行動態執行。支持提前編譯，該編譯將高級執行結構，編譯為可執行文件和生成的原始函數。所有這些執行模式，都由統一的runtime.Module 接口封裝 ，將在本文的后面部分中進行討論。
tir / transform包含用於TIR級別功能的轉換過程。許多Tir通道的目的是降維。例如，可以通過多種途徑，將多維接入到一維指針訪問，將內在函數擴展為特定於目標的內在函數，修飾函數條目，滿足運行時runtime調用約定。有一些優化過程，如簡化訪問索引和消除無效代碼。
LLVM，CUDA C和其它目標編譯器，可以在目標處理許多低級優化。將低級優化（例如寄存器分配），留給了下游編譯器，只專注於未涵蓋的優化。
搜索空間和基於學習的轉換
描述的轉換過程是確定性的，基於規則的。TVM堆棧的一個設計目標，支持針對不同硬件平台的高性能代碼優化。將需要研究盡可能多的優化選擇，包括但不限於多維張量訪問，循環切片行為，特殊的加速器內存層次結構和線程化。
很難定義做出所有選擇的試探法。將采用基於搜索和學習的方法。首先定義可以用來轉換代碼操作的集合。示例操作包括循環轉換，內聯，向量化。稱這些操作為調度原語。調度原語的集合定義了，可以對代碼進行的可能優化的搜索空間。系統搜索不同的可能調度序列，以選擇最佳調度組合。搜索過程通常以機器學習算法為原則。
搜索完成后，可以為（可能是融合的）算子，記錄最佳調度順序。編譯器可以僅查找最佳調度序列，應用於代碼。值得注意的是，此調度應用代碼布局，完全類似於基於規則的變換，能夠與傳統流程，共享相同的接口約定。
使用基於搜索的優化，處理最初的Tir函數生成問題。此模塊部分稱為AutoTVM（auto_scheduler）。隨着繼續開發TVM堆棧，希望將基於學習的轉換擴展到更多領域。
目標轉換
目標轉換階段，將IRModule轉換為相應的目標可執行格式。對於x86和ARM等后端，使用LLVM IRBuilder，構建內存中的LLVM IR。可以生成諸如CUDA C和OpenCL之類的，源代碼級語言。支持通過外部代碼生成器，將Relay函數（子圖）直接轉換為特定目標。最終代碼生成階段應盡可能輕巧。絕大部分的轉換和降維，都應在目標轉換之前進行。
提供了一個Target結構，指定編譯目標。目標翻譯階段之前的轉換，可能受到目標的影響-例如，目標的向量長度，改變向量化行為。
運行時runtime執行
TVM運行時runtime的主要目標，提供一個最小的API，使用選擇的語言（包括Python，C ++，Rust，Go，Java和JavaScript），加載和執行已編譯的工件。下面的代碼片段，顯示了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封裝編譯結果。runtime.Module包含一個GetFunction方法，用於按名稱獲取PackedFuncs。
tvm.runtime.PackedFunc，兩個生成的函數的類型清理的函數接口。runtime.PackedFunc可采用以下類型的參數並返回值：POD類型（int，float），字符串，runtime.PackedFunc，runtime.Module，runtime.NDArray以及runtime.Object的其它子類。
tvm.runtime.Module，tvm.runtime.PackedFunc，將運行時runtime模塊化的強大機制。在CUDA上獲得上述addone函數，可以使用LLVM生成主機端代碼，計算啟動參數（例如線程組的大小），從CUDAModule，調用另一個由PackedFunc支持的PackedFunc。 CUDA驅動代碼API。相同的機制可用於OpenCL內核。
上面的示例僅處理簡單的addone函數。下面的代碼段給出了使用同一接口執行端到端模型的示例：
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()
主要優點，runtime.Module和runtime.PackedFunc足以封裝算子級代碼（例如addone），及端到端模型。
總結與討論
總之，編譯流程中的關鍵數據結構為：
• IRModule：包含relay.Function和tir.PrimFunc
• runtime.Module：包含runtime.PackedFunc
編譯的大部分內容是關鍵數據結構之間的轉換。
• relay / transform和tir / transform，基於規則的確定性轉換
• auto_scheduler和autotvm，包含基於搜索的轉換
編譯流程示例，只是TVM堆棧的典型用例。將這些關鍵數據結構和轉換，開放給python和C ++ API。除了感興趣的數據結構，從numpy.ndarray更改為tvm.IRModule之外，可以像使用numpy，一樣使用TVM。以下是一些用例示例：
• 使用python API直接構造IRModule。
• 組成一組自定義的轉換（例如，自定義量化）。
• 使用TVM的python API直接操作IR。
邏輯架構組件
TVM體系結構圖
上圖顯示了項目中的主要邏輯組件。獲取有關組件及關系的信息。
tvm/support
支持模塊包含最常用的基礎構建實用代碼，例如通用競技場分配器，套接字和日志記錄。
tvm/runtime
運行時runtime是TVM堆棧的基礎。提供了加載和執行已編譯工件的機制。運行時runtime定義了一組穩定的標准C API，以與諸如Python和Rust的前端語言進行接口。
除了runtime :: PackedFunc之外，runtime :: Object是TVM運行時runtime中的，主要數據結構之一。帶有類型索引的引用計數基類，支持運行時runtime類型檢查和向下轉換。目標系統允許開發人員，向運行時runtime引入新的數據結構，例如數組，映射和新的IR數據結構。
除了部署用例之外，編譯器本身還大量使用TVM的運行時runtime機制。所有的IR數據結構都是runtime :: Object的子類，可以從Python前端直接訪問和操作。使用PackedFunc機制，將各種API開放給前端。
在運行時runtime的子目錄（例如runtime / opencl）中，定義了對不同硬件后端的運行時runtime支持。這些特定於硬件的運行時runtime模塊定義，用於設備內存分配和設備功能序列化的API。
runtime / rpc為PackedFunc實現RPC支持。可以使用RPC機制，將交叉編譯的庫，發送到遠程設備，確定執行性能的基准。rpc基礎架構，支持從廣泛的硬件后端收集數據，進行基於學習的優化。
• TVM運行系統
• 調試器
• 將VM放入TVM：Relay虛擬機
• 模塊序列化簡介
tvm/node
節點模塊在runtime :: Object的基礎上，為IR數據結構添加了其它功能。主要包括反射，序列化，結構等效和散列。
使用了節點模塊，可以通過在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節點，序列化為JSON格式，然后將其加載回。保存/存儲和檢查IR節點的能力，使編譯器更易於訪問，提供了基礎。
tvm/ir
在TVM / IR文件夾，包含跨所有IR功能變體的，統一的數據結構和接口。tvm / ir中的組件由tvm / relay和tvm / tir共享，包括
• ir模塊
• 類型
• PassContext和Pass
• 算子
功能的不同變體（例如relay.Function和tir.PrimFunc），可以共存於IRModule中。盡管這些變體可能不具有相同的內容表示，使用相同的數據結構來表示類型。使用相同的數據結構，表示這些變量的功能（類型）名稱。一旦明確定義了調用約定，統一類型系統，允許一個函數變換，調用另一個函數。這為將來的跨功能變量優化，打開了大門。
提供了一個統一的PassContext，用於配置傳遞算子，提供了通用的復合pass，執行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是表示所有系統定義的原始算子/內部函數的通用類。開發人員可以向系統注冊新的Op，以及其它屬性（例如Op是否是元素化的）。
• Pass Infrastructure
tvm/target
目標模塊包含將IRModule轉換為目標runtime.Module的，所有代碼生成器。提供了描述目標的通用Target類。
通過查詢目標中的屬性信息和注冊到每個目標id（cuda，opencl）的內置信息，可以根據目標定制編譯管道。
tvm/ir
TIR包含低級代碼表示的定義。使用tir :: PrimFunc表示，可以通過TIR傳遞轉換的函數。除IR數據結構外，tir模塊還通過公共Op注冊表，以及tir / transform中的轉換，傳遞定義了一組內置的內在函數及其屬性。
tvm/arith
此模塊與TIR緊密相關。低級代碼生成中的關鍵問題之一，分析索引的算術屬性-正則性，變量邊界以及描述迭代空間的整數集。arith模塊提供了一組進行（主要是整數）分析的工具。TIR pass可以使用這些分析，簡化和優化代碼。
tvm/te
te代表“張量表達式”。這是一個特定領域的語言模塊，允許通過編寫張量表達式，快速構建tir :: PrimFunc變體。重要的是，張量表達式本身不是可以存儲到IRModule中的自包含函數。相反，是IR的一個片段，拼接在一起，構建IRModule。
te / schedule提供了一組調度原語，控制所生成的功能。可能會將一些調度組件引入tir :: PrimFunc本身。
• InferBound Pass
• 混合前端開發人員指南
• InferBound Pass
• Hybrid Frontend Developer Guide
tvm/topi
可以針對每個用例直接通過TIR，或張量表達式（TE），構造算子。 topi（張量算子清單）提供了一組預定義的算子（在TE或TIR中），由numpy定義，在常見的深度學習工作負載中找到。提供了一組公共調度模板，在不同目標平台上，獲得高性能的實現。
tvm/relay
Relay是用於表示完整模型的高級功能性IR。在relay.transform中定義了各種優化。Relay編譯器定義了多種語言，每種語言旨在支持特定的優化樣式。值得注意的是QNN（用於導入預量化模型），VM（用於降級為動態虛擬機），內存（用於內存優化）。
• Relay IR簡介
• Relay算子策略
• 轉換布局通道
• Introduction to Relay IR
• Relay Operator Strategy
• Convert Layout Pass
tvm/autotvm
AutoTVM和AutoScheduler，都是自動進行基於搜索的代碼優化的組件。主要包括：
• Cost models和特征提取。
• 一種記錄格式，用於存儲調度基准結果，進行cost model構建。
• 一組有關代碼轉換的搜索策略。
• 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.
自動化代碼優化，仍然是活躍的研究領域。試圖對設計進行模塊化，研究人員可以通過Python綁定，快速修改組件或自己的應用算法，自定義搜索，從Python綁定中，插入算法。
• 基准性能日志格式Benchmark Performance Log Format
前端
前端將來自不同框架的模型，存放在TVM堆棧中。 tvm.relay.frontend是模型提取API的命名空間。
• TensorFlow前端