內(nèi)容一覽：本文整理自希姆計算 ToolChain 淡孝強等在 2023 Meet TVM 的分享演講，主題為「基于 TVM 的 DSA AI 編譯器的構建」，HyperAI超神經(jīng)做了不改變原意的刪減。

關鍵詞：希姆計算 DSA AI 編譯器

本文首發(fā)自 HyperAI超神經(jīng)微信公眾號~

大家好我是來自希姆計算的淡孝強，今天我將和三位同事一起來給大家分享如何在 TVM 上支持 NPU。

DSA 編譯器解決的本質(zhì)問題就是不同的模型需要部署到硬件上，利用各種抽象層級的優(yōu)化手段，使得模型盡量打滿芯片，也就是要壓縮氣泡。關于怎么去調(diào)度，Halide 描述的調(diào)度三角形是這個問題的本質(zhì)。

DSA 編譯器要解決的主要問題是什么？首先我們抽象一個 DSA 的架構，如圖所示，habana、Ascend 以及 IPU 上都是這個抽象架構的實例化。一般計算核里每個核有向量、標量以及張量的計算單元。從指令的操作和數(shù)據(jù)粒度來看，不少?DSA 可能傾向于使用相對粗粒度的指令，例如二維三維的向量和張量的指令，也有不少硬件使用細粒度的指令，例如一維的 SIMD 和 VLIW。指令間的依賴，有一些是通過顯式的接口暴露讓軟件來控制，有的是硬件自己控制。內(nèi)存是多級內(nèi)存，大多是 scratchpad memory。并行有各種粒度和維度的并行，比如 stream 并行、cluster 并行、多核并行以及不同計算部件之間的流水并行。

要支持這樣一類架構，從編譯器開發(fā)者的角度看，是從上述體系結構幾個方面對 AI 編譯器提出不同的需求，這部分后面我們會展開。

從用戶的角度看，首先要有一個穩(wěn)定和泛化的編譯器，盡量多的模型或者算子都可以編譯成功，另外，用戶希望編譯器可以提供一個可編程的界面來進行算法和算子的自定義，以確保可以獨立開展一些關鍵算法的創(chuàng)新工作。最后，對于類似我們這樣的團隊或者友商也會關注：怎么用 TVM 構建 AI 編譯器，比如怎樣管理自研和開源的 TVM 代碼，怎么搭建一套高效 CI 等等。這就是我們今天要分享的內(nèi)容，下面由我的同事來講編譯優(yōu)化的部分。

希姆計算王成科：DSA 的編譯優(yōu)化流程

本部分為希姆計算工程師王成科現(xiàn)場分享。

首先介紹一下希姆編譯實踐的整體流程概況。

針對剛才提到的架構特性，我們基于 TVM 數(shù)據(jù)結構構建了自研的優(yōu)化 pass 再加上對 TVM 的復用，組成了一個新模式實現(xiàn)：tensorturbo。

我們看到一個比較經(jīng)典的 DSA 架構，一般會提供一些高效、定制矩陣以及向量層的多核計算核心，擁有一個與之相配合的多層緩存機制，同時也會提供可并行運行的多模組執(zhí)行單元。相應的我們需要處理以下問題：

* 對數(shù)據(jù)計算進行切分、高效綁核以及定制指令的高效向量化；

* 精細管理有限的片上緩存，在不同的緩存等級上做相應的數(shù)據(jù)預??；

* 優(yōu)化多模組執(zhí)行的多級流水線，力求得到一個比較好的加速比。

這里紅色部分（上圖）顯示的是整個流程里對 TVM 復用比較高的部分，在 relay 上實現(xiàn)的圖層相關比較通用的優(yōu)化可以直接復用，另外復用程度比較高的是基于 TensorIR 和 custom LLIR 的算子實現(xiàn)部分。像我們剛才提到的跟硬件特性相關的定制優(yōu)化，則需要更多自研工作。

首先我們來看在圖層上自研的一項工作。

關注最左邊這張比較典型的計算流圖，可以看到從上到下，整體對緩存的占用及對計算的占用都在不斷減少，呈現(xiàn)倒金字塔的狀態(tài)。對于前半部分，模型規(guī)模較大時，我們需要著重解決片上緩存駐留的問題；而后半部分，在模型規(guī)模比較小的時候，需要處理計算單元利用率較低的問題。如果簡單調(diào)整模型規(guī)模，比如調(diào)整 batch size ，較小的 batch size 可以得到較低的 latency ，而相應的 throughput 會有所降低；同樣較大的 batch size 會導致 latency 較高，但是有可能提高整體 throughput。

那么我們就可以用圖調(diào)度來解決這個問題。首先，允許一個比較大的 batch size 輸入，保證全程對計算的利用率比較高，然后對整圖做一個存儲分析，加上切分和調(diào)度策略，使得模型的前半部分結果可以更好地緩存在片上，同時實現(xiàn)計算核心利用率較高的結果。實踐來看整體可以實現(xiàn) latency 和 throughput 都表現(xiàn)較好結果（詳細可以關注 OSDI 23 希姆文章：Effectively Scheduling Computational Graphs of Deep Neural Networks toward Their Domain-Specific Accelerators，6 月份可獲取鏈接查看）。

下面介紹另外一個軟流水的加速工作。

關注右上圖，實現(xiàn)了一個比較 native 的四級流水線，但明顯不是一個高效的流水線。一般高效的流水線，應該是經(jīng)過幾次迭代后，四個執(zhí)行單元都可以同步并行起來，那么這需要做一些工作，包括 L1 及 L0 上的切分、L1 上跨層的數(shù)據(jù)預取以及 L0 層級上的 double buffer 操作。通過這些工作我們可以實現(xiàn)像右下圖所展示的，加速比較高的流水線。

由此，也會引入一個新的問題，比如當多個執(zhí)行單元對緩存的同時讀寫并發(fā)數(shù)要高于當前緩存可支持的并發(fā)數(shù)時，就會產(chǎn)生競爭，這個問題會導致訪存效率成倍下降，也就是 Bank Conflict 問題。對此，我們會在編譯時靜態(tài)地對流水線進行模擬，提取沖突對象，結合 cost model 對分配地址進行交換和平移，可以極大地降低該問題的影響。

有了各種 pass 之后，可以以一個簡單的 Top-Down 方式把它們組合起來，沿著左圖中黑色流程，就得到了一個功能上可行的編譯 pipeline。但是實踐中發(fā)現(xiàn)很多問題，包括思遠提到的 pass 與 pass 之間的相互影響、缺少交互邏輯，圖層與算子之間缺少溝通邏輯等?？梢钥吹阶髨D中紅色部分指示的流程，實踐中發(fā)現(xiàn)每個路徑或者它們的組合都會導致編譯失敗。如何讓其魯棒性更強？希姆在每個可能失敗的 pass 中提供一個反饋路徑，在圖層和算子之間引入了交互邏輯，進行預分析、 prelower 操作，同時在重點部分引入一些迭代調(diào)優(yōu)機制，最終得到一個泛化性較高且調(diào)優(yōu)能力比較強的整體 pipeline 實現(xiàn)。

我們也留意到，上述工作中對數(shù)據(jù)結構的改造以及相關設計思想與目前 TVM Unity 設計有較多相似之處，我們也期待 Relax 能夠帶來更多可能性。

這里展示的是希姆在編譯流程中更加細節(jié)的 pass，從左到右就是逐層遞減的過程，其中紅色部分是對 TVM 復用比較高的，越靠近硬件特性部分會有更多的定制 pass。

下面繼續(xù)對其中的部分模塊進行詳細介紹。

希姆計算劉飛：DSA 的向量化和張量化

本部分為希姆計算工程師劉飛現(xiàn)場分享。

這個章節(jié)將展開介紹希姆向量化和張量化工作。從指令粒度考慮，指令粒度越粗，越接近 Tensor IR 的多層 loop 表達，所以向量化張量化難度越小，相反，指令粒度越細，難度也就越大，我們的 NPU 指令，支持一維/二維/三維的 tensor 數(shù)據(jù)計算。希姆也考慮過原生 TVM tensorize 過程，但考慮到 Compute Tensorize 對復雜表達能力有限，例如對 if condition 這種復雜表達式做 Tensorized 就比較困難，而且做 Tensorized 向量化后，無法 schedule。

另外當時 TensorIR Tensorize 在開發(fā)當中，不能滿足開發(fā)需求，所以希姆提供了自己的一套指令向量化流程，我們稱之為指令發(fā)射。這套流程下我們支持了大概 120 條 Tensor 指令，包括各種維度的指令等。

我們的指令流程大概分為三個模塊：

* 發(fā)射前的優(yōu)化處理。對循環(huán)軸的變換，為指令發(fā)射提供更多的發(fā)射條件和可能；

* 指令發(fā)射模塊。分析循環(huán)的結果和信息，選擇一個最優(yōu)的指令生成方式；

* 指令發(fā)射后的模塊。對指定發(fā)射處理失敗之后處理，保證在 CPU 上正確執(zhí)行。

下面是指令發(fā)射前的優(yōu)化和處理模塊，都是由一組優(yōu)化 pass 組成，其中 IfPromotion 是把阻礙循環(huán)軸發(fā)射的 if 語句盡量外提，PreProcess 是把沒有對應指令的 operator 做拆分處理，LoopShift 是對循環(huán)軸邊界為歸一化，LoopCallapse 是對連續(xù)的循環(huán)軸作盡可能的合并，LoopPartition 是做 if 相關的循環(huán)軸拆分，還有 LoopFission 是對循環(huán)內(nèi)多個 store 語句的分裂。

從這個例子可以看到，起初 IR 是不能發(fā)射任何指令的，通過優(yōu)化后，最后可以發(fā)射兩條 Tensor 指令且所有的循環(huán)軸都能夠發(fā)射指令。

再就是指令發(fā)射模塊。首先，指令發(fā)射模塊會循環(huán)分析循環(huán)中的結構，從中獲取 Optype、dtype、bufferAcess 等信息，有這些信息之后，指令識別會識別出來循環(huán)軸可能會發(fā)射哪幾種指令。因為一種 IR 結構可能對應多種 NPU 指令，所以我們會把所有可能發(fā)射的指令都識別出來，由 VectorEngine 搜索引擎去根據(jù)指令的 alignment、reshape 等一系列信息去搜索每種指令發(fā)射的可能性，最后再由 CostModel 做計算，找到最優(yōu)發(fā)射形式進行發(fā)射。

最后就是指令發(fā)射后處理模塊。主要是對指令發(fā)射失敗的 tir 做處理，保證其能在 CPU 上正確運行。還有一些特殊指令，希姆需要在算法前端打一些標記，指令發(fā)射模塊通過這些標記加上自己的 IR 分析，正確地發(fā)射相應的指令。

以上是希姆整個 DSA 張量化和向量化的流程，我們也在一些方向上做了探索，比如微內(nèi)核的方案，也是最近討論比較熱烈的方向。它的基本思想是把一個計算過程分成兩層，一層用組合微內(nèi)核的形式去拼接，另一種用搜索的方式去尋找，最終把兩層的結果做拼接，選擇一個最優(yōu)結果。這樣其優(yōu)勢是充分利用硬件資源的同時，降低搜索的復雜度，提高搜索效率。

希姆也在微內(nèi)核上做了相關探索，但考慮到微內(nèi)核方案與現(xiàn)在的解決方案相比，并沒有在性能等方面有較大提升，所以目前希姆在微內(nèi)核方向還屬于探索階段。

希姆計算袁晟：DSA 的自定義算子

本部分為希姆計算工程師袁晟現(xiàn)場分享。

首先，我們知道算子開發(fā)目前碰到了四個大問題：

* 需要支持的神經(jīng)網(wǎng)絡算子很多，進行歸類后基礎算子有 100 多個；

* 由于硬件架構不停迭代，相應指令以及算子參與的邏輯都需要進行變更；

* 性能考慮。算子融合(local memory, share memory)以及我前邊提到的圖算信息傳遞（切分等）；

* 算子需要開放給用戶，用戶只能進入軟件進行自定義算子。

我主要分成了以下三個方面介紹。首先是圖算子，圖算子是基于 relay api，把它裁剪成基礎的語言算子。

以下圖為例：

第二是元算子，所謂的元算子是基于 TVM Topi 用 compute/schedule 描述算子算法邏輯和循環(huán)變換相關邏輯。我們在開發(fā)算子時，會發(fā)現(xiàn)很多算子的 schedule 是可以復用的，基于這種情況下，希姆提供了一套類似 schedule 的模板。現(xiàn)在，我們把算子分成很多類，基于這些類，新的算子就會大量復用 schedule 模板。

接下來是一個比較復雜的算子，基于 NPU 的情況下，大家會發(fā)現(xiàn) topk、nms 等帶控制流的算法，帶很多標量計算，目前用 compute/schedule 很難描述，為解決這個問題，希姆提供一個類似 library 庫。相當于在 library 庫先編譯復雜的邏輯，然后通過跟 IR Builder 結合的方式，把整個算子的邏輯輸出。

接下來是算子的切分。對于 NPU，相對 GPU 和 CPU 情況下，TVM 每條指令都會操作連續(xù)內(nèi)存塊，同時會有 memory size 限制。同時，在這種情況下，搜索空間不大。基于這些問題，希姆提供了解決辦法，首先，會有一個候選集，把可行的解題放到候選集里，其次，對可行性進行解釋，主要考慮性能要求以及 NPU 指令限制，最后，會引入 cost function，其中會考慮算子特征以及可能用到的計算單元特征。

再接下來對算子開發(fā)比較有挑戰(zhàn)性的就是融合算子。目前面臨兩個爆炸性問題，第一不知道如何將自己的算子和其他算子組合，第二個可以看到 NPU 里有很多 memory 層級，出現(xiàn)爆炸式 memory 層級的融合。希姆 LLB 會有 shared memory 和 local memory 等融合的組合，基于這種情況下，我們也提供一個自動生成框架，先根據(jù)圖層給的調(diào)度信息，插入數(shù)據(jù)搬移操作，再根據(jù) schedule 里 master op 和 salve op 提煉 schedule info，最后根據(jù)當前指令的限制等問題做一個后處理。

最后主要展示希姆支持的算子。ONNX 算子大概是 124 個，目前支持的大概是 112 個，占比 90.3%，同時希姆有一套隨機測試，可以測試大質(zhì)數(shù)、融合組合以及一些 pattern 融合組合。

總結

本部分為希姆計算工程師淡孝強現(xiàn)場分享。

這是希姆基于 TVM 搭的 CI，這上面跑了 200 多個模型以及非常多的單元測試。MR 不搶 CI 資源的情況下，提交一個代碼需要 40 多分鐘。計算量很大，大概掛了 20 多張自研計算卡以及一些 CPU 機器。

總結，這是希姆的架構圖，如下所示：

效果來看，性能得到很大提升，另外自動生成與另一個手寫模型的團隊對標的話，基本上可以達到他們的 90% 以上。

這是希姆代碼的情況，左邊是 TVM 和自研代碼如何管理，TVM 是作為 third_party 里的數(shù)據(jù)結構來使用，希姆有自己的 source 和 python 的東西，如果我們需要對 TVM 進行更改，就在 patch 文件夾中對 TVM 進行改動。這里有三個原則：

* 大部分使用自研的 pass，也自研了 Custom module；

* patch 會限制少修改 TVM 源代碼，能 upstream 就及時 upstream；

* 定期跟 TVM 社區(qū)做同步，更新最新代碼到倉庫中。

整個代碼量也如上圖所示。

總結：

* 我們基于 TVM 端到端支持希姆一代二代芯片；

* 基于 relay 和 tir 實現(xiàn)所有的編譯優(yōu)化需求；

* 基于 tir 完成了 100+ 條向量張量指令的自動生成；

* 基于 TVM 實現(xiàn)了自定義算子方案；

* 模型一代支持 160+，二代已經(jīng)使能 20+；

* 模型性能接近手寫極致。

Q&A

Q1：我對融合算子比較感興趣，它如何跟 TVM 的 tir 結合？

A1：對于右圖，同一個算級，第一，如果算子有兩個 input 一個 output，那算子形態(tài)就有 27 種。第二，各種各樣的算子銜接時，scope 有可能是三個之一，所以我們不會假設有固定 pattern。那么如何在 TVM 上實現(xiàn)？首先根據(jù)圖層調(diào)度，決定前后 add 和中間 scope 在哪里，圖層是一個非常復雜的過程，輸出的結果是決定算子存在于哪個緩存以及可用緩存有多少。有了這個調(diào)度的結果，我們在算子層進行自動融合算子生成，比如我們根據(jù) scope 信息進行自動插入數(shù)據(jù)搬移的操作，完成數(shù)據(jù)流的構建。

schedule info 里邊和 TVM 原生的機制很類似，融合過程中需要考慮每個 member scope 所用的大小，所以這里就是 TVM 原生的東西，我們只是用了一個特別的框架，將其集成到這里，讓它自動化。

do schedule 在此基礎上，把開發(fā)者所需要的 schedule 做出來，可能也會有一些后處理。

Q2：方便透露 CostModel 更多細節(jié)嗎？cost function 是根據(jù)算子層面的 feature 還是根據(jù)硬件層面的特性結合設計的？

A1：大思路已經(jīng)在這了，首先生成一個候選集，生成過程跟 NPL 結構相關，然后會有剪枝的過程，考慮指令限制以及后邊的優(yōu)化，多核、double buffer 等，最后有一個 cost function 對其進行排序。

我們知道優(yōu)化套路本質(zhì)是如何把數(shù)據(jù)搬移隱藏在計算中，無非是對操作照此標準進行模擬，最后計算代價。

Q3：除了 TVM 支持的默認融合規(guī)則，希姆有沒有產(chǎn)生新的融合規(guī)則，比如在計算圖層結合不同硬件定制的特有融合。

A3：關于融合，實際上有兩個層次，第一，buffer，第二，loop 融合。TVM 融合方式實際上是針對后一種。希姆基本沿用你說的 TVM 融合 pattern 的套路，但是做了一些限制。