TVM 架構設計

TVM 架構設計

本文面向希望了解TVM體系結構和/或，積極參與項目開發的開發人員。

主要內容如下：             

示例編譯流程，概述了TVM將模型的高級概念，轉換為可部署模塊的步驟。            

邏輯架構組件部分，描述邏輯組件。針對每個邏輯組件，按組件的名稱進行組織。             

可以隨時查看，開發人員如何指導有用的開發技巧。             

提供了架構的一些補充視圖。檢查一個單一的端到端編譯流程，討論關鍵的數據結構和轉換。這個基於runtime的視圖，主要關注運行編譯器時，每個組件之間的交互。將檢查代碼庫的邏輯模塊及關系。設計了靜態總體視圖。

Example Compilation Flow

研究編譯器中的一個示例編譯流。下圖顯示了流程，在高層，包含幾個步驟：              

導入：前端組件將模型攝取到IRModule中，IRModule包含內部表示模型的函數集合。             

轉換：編譯器將一個IRModule，轉換成另一個功能上等價或近似等價的IRModule（如在量化的情況下）。許多轉換都是獨立於目標（后端）的。還允許target影響轉換管道的配置。             

目標轉換：編譯器將IRModule（codegen），轉換為目標指定的可執行格式。目標轉換結果封裝為runtime.Module，可以在目標runtime環境中導出、加載和執行。              

Runtime執行：用戶將runtime.Module，在支持的runtime環境中，運行編譯的函數。

Key data structures

設計和理解復雜系統的最佳方法之一，識別關鍵數據結構和操作（轉換）這些數據結構的API。一旦確定了關鍵的數據結構，可以將系統分解成邏輯組件，這些組件要么定義關鍵數據結構的集合，要么定義數據結構之間的轉換。             

IRModule是整個堆棧中，使用的主要數據結構。IRModule（中間表示模塊）包含一組函數。支持函數的兩個主要變體。             

relay::Function函數是高級函數程序表示。一個relay功能通常對應於端到端模型。可以查看relay功能作為一個計算圖，對控制流、遞歸和復雜的數據結構，有額外的支持。             

tir::PrimFunc是一個低級程序表示，包括循環嵌套選擇、多維加載/存儲、線程和向量/張量指令在內的元素。通常用來表示執行模型中（可能是融合）層的算子程序。             

在編譯過程中，一個relay函數，可以降為多個tir::PrimFunc函數，與一個調用這些tir::PrimFunc函數的頂層函數。

Transformations

已經介紹了關鍵的數據結構，談談轉換。每種轉換，都可以達到以下目的之一：             

優化：將一個程序轉換成一個等效的，可能更優化的版本。             

降低：將程序轉換為更接近目標的低級表示。             

Relay/轉換包含優化模型的過程集合。這些優化包括常見的程序優化，如常量折疊和死代碼消除，及張量計算特定的過程，如布局轉換和縮放因子折疊。             

將端到端部的Relay（例如，MobileNet），端到端部的優化（融合操作），稱這些函數段。這個過程幫助將原始問題，分為兩個子問題：             

每個子功能的編譯和優化。              

整體執行結構：需要對生成的子函數，執行一系列調用，執行整個模型。             

使用低級tir，編譯和優化每個子函數。對於特定的目標，可以直接進入目標轉換階段，使用外部代碼生成器。             

有幾種不同的方法（在Relay/后端），處理對整個執行問題的調用。對於具有已知形狀，沒有控制流的簡單模型，可以降低到將執行結構，存儲在圖形中的圖形runtime。支持用於動態執行的虛擬機后端。調度支持提前編譯，將高級執行結構，編譯成可執行的和生成的原始函數。所有這些執行模式，都被一個統一的runtime.Module接口，將在后面部分討論。             

tir/轉換包含tir級功能的轉換過程。許多tir pass的目的是降低。有一些過程，可以將多維訪問，變為一維指針訪問，將內部函數擴展為特定於目標的內部函數，及修飾函數入口，滿足runtime調用約定。有一些優化過程，如訪問索引簡化和死代碼消除。             

許多低級優化，可以在目標階段由LLVM、cudac和其它目標編譯器處理。將寄存器分配等低級優化，留給下游編譯器，只關注那些沒有覆蓋的優化。

Search-space and Learning-based Transformations

到目前為止，所描述的轉換過程是確定的和基於規則的。TVM堆棧的一個設計目標，支持針對不同硬件平台的高性能代碼優化。需要研究盡可能多的優化選擇，包括但不限於多維張量訪問、循環平鋪行為、特殊加速器內存層次結構和線程。             

很難定義一個啟發式，做出所有的選擇。相反，將采取基於搜索和學習的方法。首先定義一組，可以轉換程序的操作。示例操作包括循環轉換、內聯、矢量化。稱這些動作為調度原語。調度原語的集合，定義了一個搜索空間，可以對程序進行優化。然后系統搜索不同的可能的調度序列，選擇最佳的調度組合。搜索過程，通常由機器學習算法指導。              

一旦搜索完成，就可以記錄一個（可能是融合）算子的最佳調度序列。編譯器就可以查找最佳調度序列，應用到程序中。值得注意的是，這個調度應用程序階段與基於規則的轉換完全相同，能夠與傳統過程，共享相同的接口約定。             

使用基於搜索的優化方法，處理初始的tir函數生成問題。該模塊的這一部分，稱為AutoTVM（auto_scheduler）。隨着繼續開發TVM堆棧，將基於學習的轉換擴展到更多領域。

Target Translation

目標轉換階段，將IRModule轉換為相應的目標可執行格式。對於x86和ARM這樣的后端，使用llvmirbuilder構建內存llvmir。

可以生成源代碼級語言，如cudac和OpenCL。支持通過外部代碼生成器，將Relay函數（子圖），直接轉換為特定目標。最后的代碼生成階段，盡可能輕量級。絕大多數轉換和降低，都應該在目標轉換階段之前執行。             

還提供了一個目標結構，指定編譯目標。目標平移階段之前的變換，受到目標的影響-例如，目標的向量長度，改變矢量化行為。

Runtime Execution

TVMruntime的主要目標，提供一個最小的API，用於加載和執行編譯后的工件，語言包括Python、C++、Read、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包含按名稱獲取PackedFuncs的GetFunction方法。               

tvm.runtime.PackedFunc，兩個生成函數的類型刪除函數接口。一個runtime.PackedFunc，可以使用以下類型獲取參數和返回值：POD類型（int，float），string，runtime.PackedFunc, runtime.Module, runtime.NDArray，及其它子類runtime.Object。            

tvm.runtime.Module，tvm.runtime.PackedFunc是模塊化runtime的強大機制。例如，為了在CUDA上獲得上述addone函數，可以使用LLVM生成主機端代碼，計算啟動參數（例如線程組的大小），然后從CUDA驅動程序API支持的CUDAModule，調用另一個PackedFunc。同樣的機制，可以用於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"](tvm.gpu(0))

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"](0).asnumpy()

主要的收獲是runtime.Module 和 runtime.PackedFunc足以封裝算子級程序（如addone）和端到端模型。             

總結與討論             

總之，編譯流程中的關鍵數據結構：

IRModule: contains relay.Function and tir.PrimFunc

runtime.Module: contains runtime.PackedFunc                          

編譯的大部分內容，關鍵數據結構之間的轉換。             

Relay/轉換和tir/轉換，基於規則的確定性轉換             

auto_scheduler和autotvm包含基於搜索的轉換             

最后，編譯流程示例，只是TVM堆棧的一個典型用例。將這些關鍵數據結構和轉換，表示為Python和C++ API。可以像使用numpy一樣使用TVM，只是感興趣的數據結構，從numpy.ndarray 到 tvm.IRModule。以下是一些用例：             

使用python API直接構造IRModule。             

組成一組自定義的變換（例如自定義量化）。             

使用TVM的PythonAPI直接操作IR。

Logical Architecture Components

 

上圖顯示了項目中的主要邏輯組件。有關組件及其關系的信息，請閱讀以下部分。             

tvm/支持             

支持模塊包含架構，設計最常用的實用程序，如通用的arena分配器、socket和日志記錄。             

tvm/runtime間             

runtime是TVM堆棧的基礎。提供了加載和執行已編譯工件的機制。runtime定義了一組穩定的標准C API，與諸如Python和Rust這樣的前端語言交互。             

Object是TVM runtime中，除了runtime:：PackedFunc之外，主要數據結構之一。一個引用計數的基類，具有類型索引，支持runtime類型檢查和向下轉換。目標系統允許開發人員向runtime，引入新的數據結構，例如數組、映射和新的IR數據結構。             

除了部署用例，編譯器本身，大量使用TVM的runtime機制。所有的IR數據結構，都是runtime:：Object的子類，可以從Python前端，直接訪問和操作。使用PackedFunc機制，向前端開發各種API。             

對不同硬件后端的runtime支持，定義在runtime的子目錄中（例如Runtime/opencl）。這些特定於硬件的runtime模塊，定義了用於設備內存分配和設備函數序列化的API。             

runtime/rpc實現了對PackedFunc的rpc支持。可以使用RPC機制，將交叉編譯的庫，發送到遠程設備，對執行性能進行基准測試。Rpc架構設計，支持從通用的硬件后端收集數據，進行基於學習的優化。

• TVM Runtime System
• Debugger
• Putting the VM in TVM: The Relay Virtual Machine
• Introduction to Module Serialization

tvm/node

node模塊在IR數據結構的runtime:：Object的基礎上，添加了其它特性。主要特性包括反射、序列化、結構等價和哈希。             

有了node模塊，就可以在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共享，值得注意的是

• IRModule
• Type
• PassContext and Pass
• Op

功能的不同變體（例如，relay.Function and 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是表示所有系統定義的，基本算子/內部函數的公共類。開發人員可以向系統注冊新的算子，以及附加屬性（例如，算子是否是基本元素級的）。

• Pass Infrastructure

tvm/target

The target module contains all the code generators that translate an IRModule to a target runtime.Module. It also provides a common Target class that describes the target.

The compilation pipeline can be customized according to the target by querying the attribute information in the target and builtin information registered to each target id(cuda, opencl).

tvm/tir

TIR contains the definition of the low-level program representations. We use tir::PrimFunc to represent functions that can be transformed by TIR passes. Besides the IR data structures, the tir module also defines a set of builtin intrinsics and their attributes via the common Op registry, as well as transformation passes in tir/transform.

tvm/arith

This module is closely tied to the TIR. One of the key problems in the low-level code generation is the analysis of the indices’ arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. arith module provides a collection of tools that do (primarily integer) analysis. A TIR pass can use these analyses to simplify and optimize the code.

tvm/te

The name te stands for “tensor expression”. This is a domain-specific language module that allows us to construct tir::PrimFunc variants quickly by writing tensor expressions. Importantly, a tensor expression itself is not a self-contained function that can be stored into IRModule. Instead, it is a fragment of IR that we can stitch together to build an IRModule.

te/schedule provides a collection of scheduling primitives to control the function being generated. In the future, we might bring some of these scheduling components to the a tir::PrimFunc itself.

• InferBound Pass
• Hybrid Frontend Developer Guide

tvm/topi

While possible to construct operators directly via TIR or tensor expressions (TE) for each use case it is tedious to do so. topi (Tensor operator inventory) provides a set of pre-defined operators (in TE or TIR) defined by numpy and found in common deep learning workloads. We also provide a collection of common schedule templates to obtain performant implementations across different target platforms.

tvm/relay

Relay is the high-level functional IR used to represent full models. Various optimizations are defined in relay.transform. The Relay compiler defines multiple dialects, and each dialect is designed to support specific styles of optimization. Notable ones include QNN(for importing pre-quantized models), VM(for lowering to dynamic virtual machine), memory(for memory optimization).

• Introduction to Relay IR
• Relay Operator Strategy
• Convert Layout Pass

tvm/autotvm

AutoTVM and AutoScheduler are both components which automate search based program optimization. This is rapidly evolving and primarily consists of:

• 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.

Automated program optimization is still an active research field. As a result, we have attempted to modularize the design so that researchers may quickly modify a component or apply their own algorithms via the Python bindings, and customize the search and plugin their algorithms from the Python binding.

• Benchmark Performance Log Format

Frontends

Frontends ingest models from different frameworks into the TVM stack. tvm.relay.frontend is the namespace for model ingestion APIs.

• TensorFlow Frontend
• Security Guide

Security