內容一覽：本文整理自希姆計算 ToolChain 淡孝強等在 2023 Meet TVM 的分享演講，主題為「基於 TVM 的 DSA AI 編譯器的構建」，HyperAI超神經做了不改變原意的刪減。
關鍵詞：希姆計算 DSA AI 編譯器
大家好我是來自希姆計算的淡孝強，今天我將和三位同事一起來給大家分享如何在 TVM 上支持 NPU。
DSA 編譯器解決的本質問題就是不同的模型需要部署到硬體上，利用各種抽象層級的最佳化手段，使得模型儘量打滿晶片，也就是要壓縮氣泡。關於怎麼去排程，Halide 描述的排程三角形是這個問題的本質。
DSA 編譯器要解決的主要問題是什麼？首先我們抽象一個 DSA 的架構，如圖所示，habana、Ascend 以及 IPU 上都是這個抽象架構的實例化。一般計算核裡每個核有向量、標量以及張量的計算單元。從指令的操作和資料粒度來看，不少DSA 可能傾向於使用相對粗粒度的指令，例如二維三維的向量和張量的指令，也有不少硬體使用細粒度的指令，例如一維的 SIMD 和 VLIW。指令間的依賴，有一些是通過顯式的接口暴露讓軟體來控制，有的是硬體自己控制。記憶體是多級記憶體，大多是 scratchpad memory。並行有各種粒度和維度的並行，比如 stream 並行、cluster 並行、多核並行以及不同計算部件之間的流水並行。
要支持這樣一類架構，從編譯器開發者的角度看，是從上述體系結構幾個方面對 AI 編譯器提出不同的需求，這部分後面我們會展開。
從使用者的角度看，首先要有一個穩定和泛化的編譯器，儘量多的模型或者運算元都可以編譯成功，另外，使用者希望編譯器可以提供一個可程式設計的界面來進行演算法和運算元的自定義，以確保可以獨立開展一些關鍵演算法的創新工作。最後，對於類似我們這樣的團隊或者友商也會關注：怎麼用 TVM 構建 AI 編譯器，比如怎樣管理自研和開源的 TVM 程式碼，怎麼搭建一套高效 CI 等等。這就是我們今天要分享的內容，下面由我的同事來講編譯最佳化的部分。
希姆計算王成科：DSA 的編譯最佳化流程
本部分為希姆計算工程師王成科現場分享。
首先介紹一下希姆編譯實踐的整體流程概況。
針對剛才提到的架構特性，我們基於 TVM 資料結構構建了自研的最佳化 pass 再加上對 TVM 的複用，組成了一個新模式實現：tensorturbo。
我們看到一個比較經典的 DSA 架構，一般會提供一些高效、定製矩陣以及向量層的多核計算核心，擁有一個與之相配合的多層快取機制，同時也會提供可並行運行的多模組執行單元。相應的我們需要處理以下問題：
* 對資料計算進行切分、高效綁核以及定製指令的高效向量化；
* 精細管理有限的片上快取，在不同的快取等級上做相應的資料預取；
* 最佳化多模組執行的多級流水線，力求得到一個比較好的加速比。
這裡紅色部分（上圖）顯示的是整個流程裡對 TVM 複用比較高的部分，在 relay 上實現的圖層相關比較通用的最佳化可以直接複用，另外複用程度比較高的是基於 TensorIR 和 custom LLIR 的運算元實現部分。像我們剛才提到的跟硬體特性相關的定製最佳化，則需要更多自研工作。
首先我們來看在圖層上自研的一項工作。
關注最左邊這張比較典型的計算流圖，可以看到從上到下，整體對快取的佔用及對計算的佔用都在不斷減少，呈現倒金字塔的狀態。對於前半部分，模型規模較大時，我們需要著重解決片上快取駐留的問題；而後半部分，在模型規模比較小的時候，需要處理計算單元利用率較低的問題。如果簡單調整模型規模，比如調整 batch size ，較小的 batch size 可以得到較低的 latency ，而相應的 throughput 會有所降低；同樣較大的 batch size 會導致 latency 較高，但是有可能提高整體 throughput。
那麼我們就可以用圖排程來解決這個問題。首先，允許一個比較大的 batch size 輸入，保證全程對計算的利用率比較高，然後對整圖做一個儲存分析，加上切分和排程策略，使得模型的前半部分結果可以更好地快取在片上，同時實現計算核心利用率較高的結果。實踐來看整體可以實現 latency 和 throughput 都表現較好結果（詳細可以關注 OSDI 23 希姆文章：Effectively Scheduling Computational Graphs of Deep Neural Networks toward Their Domain-Specific Accelerators，6 月份可獲取連結查看）。
下面介紹另外一個軟流水的加速工作。
關注右上圖，實現了一個比較 native 的四級流水線，但明顯不是一個高效的流水線。一般高效的流水線，應該是經過幾次迭代後，四個執行單元都可以同步並行起來，那麼這需要做一些工作，包括 L1 及 L0 上的切分、L1 上跨層的資料預取以及 L0 層級上的 double buffer 操作。通過這些工作我們可以實現像右下圖所展示的，加速比較高的流水線。
由此，也會引入一個新的問題，比如當多個執行單元對快取的同時讀寫併發數要高於當前快取可支持的併發數時，就會產生競爭，這個問題會導致訪存效率成倍下降，也就是 Bank Conflict 問題。對此，我們會在編譯時靜態地對流水線進行模擬，提取衝突對象，結合 cost model 對分配地址進行交換和平移，可以極大地降低該問題的影響。
有了各種 pass 之後，可以以一個簡單的 Top-Down 方式把它們組合起來，沿著左圖中黑色流程，就得到了一個功能上可行的編譯 pipeline。但是實踐中發現很多問題，包括思遠提到的 pass 與 pass 之間的相互影響、缺少互動邏輯，圖層與運算元之間缺少溝通邏輯等。可以看到左圖中紅色部分指示的流程，實踐中發現每個路徑或者它們的組合都會導致編譯失敗。如何讓其魯棒性更強？希姆在每個可能失敗的 pass 中提供一個反饋路徑，在圖層和運算元之間引入了互動邏輯，進行預分析、 prelower 操作，同時在重點部分引入一些迭代調優機制，最終得到一個泛化性較高且調優能力比較強的整體 pipeline 實現。
我們也留意到，上述工作中對資料結構的改造以及相關設計思想與目前 TVM Unity 設計有較多相似之處，我們也期待 Relax 能夠帶來更多可能性。
這裡展示的是希姆在編譯流程中更加細節的 pass，從左到右就是逐層遞減的過程，其中紅色部分是對 TVM 複用比較高的，越靠近硬體特性部分會有更多的定製 pass。
下面繼續對其中的部分模組進行詳細介紹。
希姆計算劉飛：DSA 的向量化和張量化
本部分為希姆計算工程師劉飛現場分享。
這個章節將展開介紹希姆向量化和張量化工作。從指令粒度考慮，指令粒度越粗，越接近 Tensor IR 的多層 loop 表達，所以向量化張量化難度越小，相反，指令粒度越細，難度也就越大，我們的 NPU 指令，支持一維/二維/三維的 tensor 資料計算。希姆也考慮過原生 TVM tensorize 過程，但考慮到 Compute Tensorize 對複雜表達能力有限，例如對 if condition 這種複雜表達式做 Tensorized 就比較困難，而且做 Tensorized 向量化後，無法 schedule。
另外當時 TensorIR Tensorize 在開發當中，不能滿足開發需求，所以希姆提供了自己的一套指令向量化流程，我們稱之為指令發射。這套流程下我們支持了大概 120 條 Tensor 指令，包括各種維度的指令等。
我們的指令流程大概分為三個模組：
* 發射前的最佳化處理。對循環軸的變換，為指令發射提供更多的發射條件和可能；
* 指令發射模組。分析循環的結果和資訊，選擇一個最優的指令生成方式；
* 指令發射後的模組。對指定發射處理失敗之後處理，保證在 CPU 上正確執行。
下面是指令發射前的最佳化和處理模組，都是由一組最佳化 pass 組成，其中 IfPromotion 是把阻礙循環軸發射的 if 語句儘量外提，PreProcess 是把沒有對應指令的 operator 做拆分處理，LoopShift 是對循環軸邊界為歸一化，LoopCallapse 是對連續的循環軸作儘可能的合併，LoopPartition 是做 if 相關的循環軸拆分，還有 LoopFission 是對循環內多個 store 語句的分裂。
從這個例子可以看到，起初 IR 是不能發射任何指令的，通過最佳化後，最後可以發射兩條 Tensor 指令且所有的循環軸都能夠發射指令。
再就是指令發射模組。首先，指令發射模組會循環分析循環中的結構，從中獲取 Optype、dtype、bufferAcess 等資訊，有這些資訊之後，指令識別會識別出來循環軸可能會發射哪幾種指令。因為一種 IR 結構可能對應多種 NPU 指令，所以我們會把所有可能發射的指令都識別出來，由 VectorEngine 搜尋引擎去根據指令的 alignment、reshape 等一系列資訊去搜尋每種指令發射的可能性，最後再由 CostModel 做計算，找到最優發射形式進行發射。
最後就是指令發射後處理模組。主要是對指令發射失敗的 tir 做處理，保證其能在 CPU 上正確運行。還有一些特殊指令，希姆需要在演算法前端打一些標記，指令發射模組通過這些標記加上自己的 IR 分析，正確地發射相應的指令。
以上是希姆整個 DSA 張量化和向量化的流程，我們也在一些方向上做了探索，比如微核心的方案，也是最近討論比較熱烈的方向。它的基本思想是把一個計算過程分成兩層，一層用組合微核心的形式去拼接，另一種用搜尋的方式去尋找，最終把兩層的結果做拼接，選擇一個最優結果。這樣其優勢是充分利用硬體資源的同時，降低搜尋的複雜度，提高搜尋效率。
希姆也在微核心上做了相關探索，但考慮到微核心方案與現在的解決方案相比，並沒有在性能等方面有較大提升，所以目前希姆在微核心方向還屬於探索階段。
希姆計算袁晟：DSA 的自定義運算元
本部分為希姆計算工程師袁晟現場分享。
首先，我們知道運算元開發目前碰到了四個大問題：
* 需要支持的神經網路運算元很多，進行歸類後基礎運算元有 100 多個；
* 由於硬體架構不停迭代，相應指令以及運算元參與的邏輯都需要進行變更；
* 性能考慮。運算元融合(local memory, share memory)以及我前邊提到的圖算資訊傳遞（切分等）；
* 運算元需要開放給使用者，使用者只能進入軟體進行自定義運算元。
我主要分成了以下三個方面介紹。首先是圖運算元，圖運算元是基於 relay api，把它裁剪成基礎的語言運算元。
以下圖為例：
第二是元運算元，所謂的元運算元是基於 TVM Topi 用 compute/schedule 描述運算元演算法邏輯和循環變換相關邏輯。我們在開發運算元時，會發現很多運算元的 schedule 是可以複用的，基於這種情況下，希姆提供了一套類似 schedule 的模板。現在，我們把運算元分成很多類，基於這些類，新的運算元就會大量複用 schedule 模板。
接下來是一個比較複雜的運算元，基於 NPU 的情況下，大家會發現 topk、nms 等帶控制流的演算法，帶很多標量計算，目前用 compute/schedule 很難描述，為解決這個問題，希姆提供一個類似 library 庫。相當於在 library 庫先編譯複雜的邏輯，然後通過跟 IR Builder 結合的方式，把整個運算元的邏輯輸出。
接下來是運算元的切分。對於 NPU，相對 GPU 和 CPU 情況下，TVM 每條指令都會操作連續記憶體塊，同時會有 memory size 限制。同時，在這種情況下，搜尋空間不大。基於這些問題，希姆提供了解決辦法，首先，會有一個候選集，把可行的解題放到候選集裡，其次，對可行性進行解釋，主要考慮性能要求以及 NPU 指令限制，最後，會引入 cost function，其中會考慮運算元特徵以及可能用到的計算單元特徵。
再接下來對運算元開發比較有挑戰性的就是融合運算元。目前面臨兩個爆炸性問題，第一不知道如何將自己的運算元和其他運算元組合，第二個可以看到 NPU 裡有很多 memory 層級，出現爆炸式 memory 層級的融合。希姆 LLB 會有 shared memory 和 local memory 等融合的組合，基於這種情況下，我們也提供一個自動生成框架，先根據圖層給的排程資訊，插入資料搬移操作，再根據 schedule 裡 master op 和 salve op 提煉 schedule info，最後根據當前指令的限制等問題做一個後處理。
最後主要展示希姆支持的運算元。ONNX 運算元大概是 124 個，目前支持的大概是 112 個，佔比 90.3%，同時希姆有一套隨機測試，可以測試大質數、融合組合以及一些 pattern 融合組合。
總結
本部分為希姆計算工程師淡孝強現場分享。
這是希姆基於 TVM 搭的 CI，這上面跑了 200 多個模型以及非常多的單元測試。MR 不搶 CI 資源的情況下，提交一個程式碼需要 40 多分鐘。計算量很大，大概掛了 20 多張自研計算卡以及一些 CPU 機器。
總結，這是希姆的架構圖，如下所示：
效果來看，性能得到很大提升，另外自動生成與另一個手寫模型的團隊對標的話，基本上可以達到他們的 90% 以上。
這是希姆程式碼的情況，左邊是 TVM 和自研程式碼如何管理，TVM 是作為 third_party 裡的資料結構來使用，希姆有自己的 source 和 python 的東西，如果我們需要對 TVM 進行更改，就在 patch 檔案夾中對 TVM 進行改動。這裡有三個原則：
* 大部分使用自研的 pass，也自研了 Custom module；
* patch 會限制少修改 TVM 源程式碼，能 upstream 就及時 upstream；
* 定期跟 TVM 社區做同步，更新最新程式碼到倉庫中。
整個程式碼量也如上圖所示。
總結：
* 我們基於 TVM 端到端支持希姆一代二代晶片；
* 基於 relay 和 tir 實現所有的編譯最佳化需求；
* 基於 tir 完成了 100+ 條向量張量指令的自動生成；
* 基於 TVM 實現了自定義運算元方案；
* 模型一代支持 160+，二代已經使能 20+；
* 模型性能接近手寫極致。
Q&A
Q1：我對融合運算元比較感興趣，它如何跟 TVM 的 tir 結合？
A1：對於右圖，同一個算級，第一，如果運算元有兩個 input 一個 output，那運算元形態就有 27 種。第二，各種各樣的運算元銜接時，scope 有可能是三個之一，所以我們不會假設有固定 pattern。那麼如何在 TVM 上實現？首先根據圖層排程，決定前後 add 和中間 scope 在哪裡，圖層是一個非常複雜的過程，輸出的結果是決定運算元存在於哪個快取以及可用快取有多少。有了這個排程的結果，我們在運算元層進行自動融合運算元生成，比如我們根據 scope 資訊進行自動插入資料搬移的操作，完成資料流的構建。
schedule info 裡邊和 TVM 原生的機制很類似，融合過程中需要考慮每個 member scope 所用的大小，所以這裡就是 TVM 原生的東西，我們只是用了一個特別的框架，將其集成到這裡，讓它自動化。
do schedule 在此基礎上，把開發者所需要的 schedule 做出來，可能也會有一些後處理。
Q2：方便透露 CostModel 更多細節嗎？cost function 是根據運算元層面的 feature 還是根據硬體層面的特性結合設計的？
A1：大思路已經在這了，首先生成一個候選集，生成過程跟 NPL 結構相關，然後會有剪枝的過程，考慮指令限制以及後邊的最佳化，多核、double buffer 等，最後有一個 cost function 對其進行排序。
我們知道最佳化套路本質是如何把資料搬移隱藏在計算中，無非是對操作照此標準進行模擬，最後計算代價。
Q3：除了 TVM 支持的默認融合規則，希姆有沒有產生新的融合規則，比如在計算圖層結合不同硬體定製的特有融合。
A3：關於融合，實際上有兩個層次，第一，buffer，第二，loop 融合。TVM 融合方式實際上是針對後一種。希姆基本沿用你說的 TVM 融合 pattern 的套路，但是做了一些限制。
溫馨提示
獲取 PPT：公眾號後臺回覆TVM 上海，獲取本場分享的完整 PPT。
