內(nèi)容一覽：本文整理自 Apache TVM PMC、上海交通大學(xué)博士生馮思遠(yuǎn)在 2023 Meet TVM 上海站的演講分享，主題為「TVM 與機(jī)器學(xué)習(xí)編譯發(fā)展」，HyperAI超神經(jīng)做了不改變?cè)獾膭h減整理。

關(guān)鍵詞：TVM 2023 Meet TVM 機(jī)器學(xué)習(xí)編譯

下午好，歡迎大家今天來參加 2023 Meet TVM。作為 Apache TVM PMC，由我來給大家做關(guān)于 TVM 的發(fā)展以及 TVM 未來 Unity 框架的分享。

Apache TVM Evolution

首先為什么會(huì)有 MLC (Machine Learning Compilation)？隨著 AI 模型的不斷擴(kuò)展，實(shí)際生產(chǎn)應(yīng)用中會(huì)有涌現(xiàn)更多的需求，很多應(yīng)用的第一層 AI Application（如圖所示）是共用的，包括 ResNet、BERT、Stable Diffusion 等模型。

第二層 Scenario 就不一樣了，開發(fā)者需要在不同的場(chǎng)景去部署這些模型，最開始是云計(jì)算、高性能計(jì)算，需要用 GPU 進(jìn)行加速。隨著 AI 領(lǐng)域的加速，最重要的任務(wù)就是把它帶進(jìn)千家萬戶，也就是 Personal PC 個(gè)人電腦、Mobile 手機(jī)以及 Edge 設(shè)備。

然而不同的場(chǎng)景有不同的需求，包括降低成本、提升性能。像 Out of the Box 需要保證用戶打開網(wǎng)頁(yè)或者下載一個(gè)應(yīng)用，可以立馬使用，手機(jī)需要省電，Edge 需要在沒有 OS 的硬件上運(yùn)行起來，有些時(shí)候還需要在低功耗、低算力的芯片上把技術(shù)跑起來。這些就是大家在不同應(yīng)用上遇到的困難，怎么解決？

MLC 領(lǐng)域?qū)Υ耸怯泄沧R(shí)的，即 Muli-Level IR Design。核心有三層，第一層 Graph-Level IR，中間層 Tensor-Level IR 以及下一層 Hardware-Level IR。這幾層是必須的，因?yàn)槟Ｐ褪且粋€(gè) Graph，中間層是 Tensor-Level IR，MLC 的核心就是優(yōu)化 Tensor Computing。底下這兩層 Hardware-Level IR 和 Hardware 是相互綁定的，就是說 TVM 不會(huì)去涉及直接生成匯編指令這一層，因?yàn)橹虚g會(huì)有一些更細(xì)的優(yōu)化技巧，這一層交由廠商或者編譯器去解決。

ML Compiler 在設(shè)計(jì)之初有一下目標(biāo)：

• Dependency Minimization

第一點(diǎn)，最小化依賴部署。為什么現(xiàn)在 AI 應(yīng)用沒有真正落地，正是因?yàn)椴渴鸬拈T檻太高。運(yùn)行過 ChatGPT 的人比 Stable Diffusion 多并不是因?yàn)?Stable Diffusion 不夠厲害，而是因?yàn)?ChatGPT 提供了一個(gè)開箱可用的環(huán)境。在我看來，使用 Stable Diffusion，你需要先從 GitHub 上下載一個(gè)模型，然后開一個(gè) GPU 服務(wù)器把它部署起來，但是 ChatGPT 則是開箱即用。開箱即用關(guān)鍵的一點(diǎn)就是最小化依賴，所有人、所有環(huán)境上都可以用。

• Various Hardware Support

第二點(diǎn)，能夠支持不同硬件。多樣化硬件部署在發(fā)展早期不是一個(gè)最重要的題，但隨著國(guó)內(nèi)外 AI 芯片的發(fā)展，它會(huì)變得越來越重要，尤其是國(guó)內(nèi)目前的環(huán)境以及國(guó)內(nèi)芯片公司的現(xiàn)狀，需要我們對(duì)各個(gè)硬件都有一個(gè)很好的支持。

• Compilation Optimization

第三點(diǎn)，通用編譯優(yōu)化。通過上幾層IR的編譯可以優(yōu)化性能，包括提升運(yùn)行效率、減少內(nèi)存的占用等。

現(xiàn)在大家大多將邊編譯邊優(yōu)化作為最重要的點(diǎn)，但對(duì)整個(gè)社區(qū)來說，前兩點(diǎn)是關(guān)鍵的。因?yàn)檫@是從編譯器角度來看的，并且這兩點(diǎn)都是從零到一的突破，優(yōu)化性能往往是錦上添花。

回到演講主題，我將 TVM 的發(fā)展劃分為四個(gè)階段，僅代表個(gè)人觀點(diǎn)。

TensorIR Abstraction

Stage 1：在這個(gè)階段，TVM 在 CPU、 GPU 上做推理的優(yōu)化及加速，GPU 特指 SIMT 的硬件部分。這個(gè)階段很多云計(jì)算廠商因?yàn)榘l(fā)現(xiàn)在 CPU 和 GPU 上都能加速開始用 TVM。為什么？我前面講到 CPU 和 GPU 沒有 Tensorization Support。TVM 初代 TE Schedule 基于 Halide，沒有良好的 Tensorization Support ，所以 TVM 后續(xù)發(fā)展以 Halide 發(fā)展的技術(shù)路線，包括 Auto TVM 和 Ansor，都對(duì) Tensorization 的支持不友好。

首先來看硬件發(fā)展過程，從 CPU 到 GPU 是 2015、2016 年左右，從 GPU 到 TPU 是 2019 年左右。為了做 Tensorization 支持，TVM 先分析了 Tensorized Programs 特征。

• Optimized loop nests with thread binding

第一，需要一個(gè) loop testing，這是所有 Tensorized Program 必需的，底下有 Multi-dimensional data load，這和 CMT、CPU 不一樣，它以張量而不是以標(biāo)量為單位儲(chǔ)存和計(jì)算的。

• Multi-dimensional data load/store into specialized memory buffer

第二，它儲(chǔ)存在一個(gè)特殊 memory buffer 里。

• Opaque tensorized computation body 16x16 matrix multiplication

第三，會(huì)有一個(gè)硬件池允許計(jì)算。以下的 Tensor Primitive 為例子，計(jì)算 16* 16 的矩陣乘法，這個(gè)計(jì)算不再會(huì)表達(dá)成以標(biāo)量組合的計(jì)算模式，而是會(huì)以一個(gè)指令去計(jì)算成賬號(hào)單元。

依據(jù)以上三點(diǎn) Tensorized Program 的定性分析，TVM 引入了 Computational Block。Block 是一個(gè)計(jì)算單元，其最外層有 nesting，中間有 iteration iterator、dependency relations，最底下是 body，其概念就是把內(nèi)外層的計(jì)算分離，即 Isolate the internal computation tensorized computation。

Stage 2：在這個(gè)階段，TVM 做 Auto-Tensorrization。具體如何實(shí)現(xiàn)的，這里一個(gè)例子詳細(xì)展開。

Auto-Tensorization

輸入 Input Program 和 Tensor Intrinsic，結(jié)果表明：TensorIR 和 TensorRT 在 GPU 上基本上打平，而在一些標(biāo)準(zhǔn)模型性能并不是很好。因?yàn)闃?biāo)準(zhǔn)模型是 ML Perf 的標(biāo)準(zhǔn)指標(biāo)，NVIDIA 工程師會(huì)花很多時(shí)間去做。在標(biāo)準(zhǔn)模型上表現(xiàn)優(yōu)于 TensorRT 是比較少見的，這也是相當(dāng)于打贏了業(yè)界最先進(jìn)的技術(shù)。

這是 ARM 自研 CPU 上的性能比較，TensorIR 和 ArmComputeLib 在 End-to-End 上可以比 Ansor 和 PyTorch 快約 2 倍。性能不是最關(guān)鍵的，Auto-Tensorization 思想才最為核心。

Stage 3：An End-to-End ML Compiler for Tensorized Hardware。這個(gè)階段可以將它扔到 GPU 或者已經(jīng)支持的加速器芯片上，接下來有自動(dòng)調(diào)優(yōu)、模型導(dǎo)入，是一套自洽的系統(tǒng)。這個(gè)階段 TVM 的核心是 End-to-End，可以開發(fā)一個(gè)直接使用的模型，但自定義難度很高。

接下來，關(guān)于 Relax、Unity 的發(fā)展和思考我會(huì)講得慢一點(diǎn)，原因是：

• 個(gè)人認(rèn)為 Relax、Unity 比較重要；

• 尚處于實(shí)驗(yàn)階段，很多東西只有思想，缺少 End-to-End demo 及完整代碼。

  
  

Apache TVM Stack 的局限性：

• Huge Gap between Relay and TIR。TVM 最大問題是 relay 到 TIR 的編譯范式太陡峭；

• Fixed Pipeline for Most Hardware。TVM 標(biāo)準(zhǔn)的流程是 Relay 到 TIR 到？編譯下去，實(shí)際上很多硬件要么只支持 BYOC，要么想通過 BYOC+TIR，Relay 支持的并不好，Either TIR or Library。以 GPU 加速為例，給 Relay 底層是固定的東西，要么寫 CUDA 做 Auto Tuning，要么給 TensorRT 走 BYOC，要不走 cuBLAS 調(diào)第三方庫(kù)。雖然有諸多選擇，但都是一個(gè)二選一的問題。這個(gè)問題影響比較大，而且在 Relay 上不好解決。

解決方案：TVM Unity。

Apache TVM Unity

把 Relax 和 TIR 這兩層 IR 當(dāng)成一個(gè)整體，融合成了 Graph-TIR 編程范式。融合的形式：以最簡(jiǎn)單的 Linear 模型為例，這種情況下整個(gè) lR 是可控可編程的，用 Graph-TIR 這一層的語(yǔ)言解決了 lower 過于陡峭的問題，高層算子可以逐步修改，甚至可以自行改成任何 BYOC 或 Function Calls。

1. 支持自定義解決 building Pipeline 的問題。原本 TVM Pipeline 是 Relay 額外的 TIR，然后 TIR 做 Tuning，Tuning 完成后傳遞給 LLVM 或 CUDA，這是一套 fix 的 building Pipeline，現(xiàn)在 Pipeline 發(fā)生了改變。

2. Choose the Best Perf among All Paths。開發(fā)者可以在 library 或 TIR 里做選擇，調(diào)任意東西。這是 Unity 解決的最重要的問題，社區(qū)認(rèn)為這是 unified ML compilation solution。

Misunderstanding

• TVM 和 MLIR 是競(jìng)爭(zhēng)關(guān)系

實(shí)際上，TVM 和 MLIR 沒有明確的、同層次的競(jìng)爭(zhēng)關(guān)系，TVM 專注于 MLC 機(jī)器學(xué)習(xí)編譯，MLIR 強(qiáng)調(diào) Multi-Level，也可以利用其特性來做 ML 編譯。開發(fā)者使用 MLIR 做機(jī)器學(xué)習(xí)編譯，一方面是因?yàn)?MLIR 與PyTorch等框架有原生對(duì)接，另一方面是因?yàn)樵?Unity 之前， TVM 的定制化能力比較弱。

• TVM = Inference Engine for CPU/GPU

TVM 從來不是一個(gè) inference engine，它能夠做編譯， 開發(fā)者可以用其來做推理加速。TVM 是 compiler infrastructure，但不是 inference engine，「TVM 只能用來加速」這種想法是錯(cuò)誤的。TVM 之所以能用來加速，本質(zhì)原因是 Compiler，比 PyTorch 等 Eager Mode 的執(zhí)行方式要快。

• TVM = Auto Tuning

Relax 出現(xiàn)之前，大家對(duì) TVM 的第一反應(yīng)是，能夠通過 Auto Tuning 獲得更好的性能。TVM 接下來的發(fā)展方向是淡化這個(gè)概念。TVM 提供了各種途徑以達(dá)到更好的性能、定制化整套編譯流程。TVM Unity 要做的就是提供一套架構(gòu)，把各種優(yōu)勢(shì)組合起來。

Next Step

下一步 TVM 要做 Cross-Layer Machine Learning Compiler Infrastructure 以及致力于成為 customizable Building Pipeline for Different Backends，在不同的硬件上支持可定制化，這一點(diǎn)需要結(jié)合各種方式、集各家所長(zhǎng)。

Q & A

Q 1：TVM 有沒有計(jì)劃優(yōu)化大模型？

A 1：在大模型方面，我們有一些初步的想法。目前 TVM 已經(jīng)開始做分布式推理以及一些簡(jiǎn)單訓(xùn)練，但距離真正落地還需要一段時(shí)間。

Q 2：Relax 后續(xù)在 Dynamic Shape 方面會(huì)有哪些支持和演進(jìn)？

A 2：Relay vm 支持 Dynamic shape，但沒有生成 Symbolic Deduction，例如 nmk 的矩陣乘法其輸出是 n 和 m，但在 Relay 的表達(dá)方式是 3 個(gè) nmk 統(tǒng)稱為 Any，即未知 dimension，其輸出也是未知 dimension。Relay VM 能運(yùn)行這些tasks，但是在編譯階段會(huì)丟失一些信息，所以 Relax 解決了這些問題，這就是 Relax 針對(duì) Relay 在 Dynamic Shape 上的改進(jìn)。

Q 3：TVM 的優(yōu)化和 Device 的優(yōu)化的結(jié)合問題。用 Graph 直接生成指令 的話，Device 的優(yōu)化中 TE、TIR 好像不太用到。用 BYOC 的話，好像也是跳過 TE、TIR。分享中提到 Relax 可能會(huì)有一些定制化，好像可以解決這一問題。

A 3：實(shí)際上已經(jīng)有很多硬件廠商走通了 TIR 這條路，還有一些廠商沒有關(guān)注相關(guān)的技術(shù)，選擇的還是 BYOC 方法。BYOC 不是嚴(yán)格意義上的 Compilation，在 Building Pattern 上是有限制的。綜合來說，不是企業(yè)用不上社區(qū)技術(shù)，而是依據(jù)自身情況做出了不一樣的選擇。

Q 4：TVM Unity 的出現(xiàn)是否會(huì)涉及到較高的遷移成本？從 TVM PMC 的角度來看，如何幫助用戶順暢過渡到 TVM Unity？

A 4：TVM 社區(qū)并沒有放棄 Relay，只是新增了 Relax 這個(gè)選項(xiàng)，因此老版本會(huì)持續(xù)演進(jìn)，但是為了使用一些新功能，可能需要做一些代碼及版本的遷移。

在Relax完整發(fā)布后，社區(qū)會(huì)提供的遷移教程和一定的工具支持，支持 Relay 模型直接導(dǎo)入到 Relax。但是，基于 Relay 開發(fā)的定制版本遷移到 Relax，還是需要一定的工作量，這種工作量對(duì)于公司十幾人的團(tuán)隊(duì)，大概需要一個(gè)月。

Q 5：TensorIR 在 Tensor 這塊相比之前有很大的進(jìn)展，但我注意到 TensorIR 主要還是針對(duì)編程模型如 SIMT、SIMD 這些成熟的編程方式做的，TensorIR 在現(xiàn)在很多的新的 AI 芯片、編程模型這里有沒有一些進(jìn)展?

A 5：從社區(qū)角度來說，TensorIR 之所以做 SIMT 模型，是因?yàn)楝F(xiàn)在只能利用到 SIMT 硬件，很多廠商的硬件和指令集都不開源，能夠接觸到大廠的硬件基本只有 CPU、GPU 和部分手機(jī) SoC，其他的廠商社區(qū)基本接觸不到，所以不能根據(jù)它們的編程模型做。此外，即使社區(qū)和企業(yè)合作做出來了類似層級(jí)的 TIR ，也無法開源，這是出于商業(yè)運(yùn)營(yíng)的考量。

以上為馮思遠(yuǎn)在 2023 Meet TVM 上海站的演講整理內(nèi)容。接下來，本次活動(dòng)其他嘉賓分享的詳細(xì)內(nèi)容也將陸續(xù)在「HyperAI超神經(jīng)」公眾號(hào)中發(fā)布推出，歡迎持續(xù)關(guān)注！

獲取 PPT：關(guān)注微信公眾號(hào)「HyperAI超神經(jīng)」，后臺(tái)回復(fù)關(guān)鍵字?TVM 上海，獲取完整 PPT。

TVM 中文文檔：https://tvm.hyper.ai/

GitHub 地址：https://github.com/apache/tvm