點(diǎn)擊下方卡片，關(guān)注“新機(jī)器視覺”公眾號(hào)
視覺/圖像重磅干貨，第一時(shí)間送達(dá)
陳天奇是機(jī)器學(xué)習(xí)領(lǐng)域著名的青年華人學(xué)者之一，本科畢業(yè)于上海交通大學(xué)ACM班，博士畢業(yè)于華盛頓大學(xué)計(jì)算機(jī)系，研究方向?yàn)榇笠?guī)模機(jī)器學(xué)習(xí)。在本文中，陳天奇回答了目前深度學(xué)習(xí)編譯技術(shù)的瓶頸在哪里，下一代技術(shù)是什么？
2021年12月16日，第四次TVMCon如期在線舉行。從大約四年多前行業(yè)的起點(diǎn)開始，深度學(xué)習(xí)編譯技術(shù)也從初期萌芽階段到了現(xiàn)在行業(yè)廣泛參與的狀況。我們也看到了許多關(guān)于編譯和優(yōu)化技術(shù)有趣的討論。深度學(xué)習(xí)編譯技術(shù)生態(tài)的蓬勃發(fā)展也使得我們有了許多新的思考。TVM作為最早打通的深度學(xué)習(xí)編譯技術(shù)框架已經(jīng)可以給大家提供不少價(jià)值。但是就好像深度學(xué)習(xí)框架本身會(huì)經(jīng)歷從從第一代（caffe）到下兩代（TF, pytorch）的變化一樣，我們非常清楚深度學(xué)習(xí)編譯技術(shù)本身也需要經(jīng)歷幾代的演化。因此從兩年前開始我們就開始問這樣一個(gè)問題：“目前深度學(xué)習(xí)編譯技術(shù)的瓶頸在哪里，下一代技術(shù)是什么“。比較幸運(yùn)的是，作為最早打通編譯流程的架構(gòu)我們可以比較早地直接觀察到只有嘗試整合才可以看到的經(jīng)驗(yàn)。而目前開始繁榮的生態(tài)本身 (MLIR的各種dialects，XLA, ONNX, TorchScript) 也開始給我們不少可以學(xué)習(xí)參考的目標(biāo)。本文是對(duì)過去兩年全面深度學(xué)習(xí)編譯和硬件加速生態(tài)的總結(jié)思考，也是我們對(duì)于深度學(xué)習(xí)新一代編譯技術(shù)的技術(shù)展望，希望對(duì)大家有一些參考價(jià)值。本文的內(nèi)容大部分來自于今年我們?cè)赥VMCon的主題報(bào)告。現(xiàn)狀是什么：當(dāng)前深度學(xué)習(xí)編譯解決方案和瓶頸
四類抽象
當(dāng)然深度學(xué)習(xí)編譯加速生態(tài)已經(jīng)從萌芽階段到開始繁榮生長(zhǎng)的階段。但是拋開具體實(shí)現(xiàn)而言，現(xiàn)在深度學(xué)習(xí)編譯生態(tài)圍繞著四類抽象展開:計(jì)算圖表示(computational graph)：計(jì)算圖可以把深度學(xué)習(xí)程序表示成DAG，然后進(jìn)行類似于算子融合，改寫，并行等高級(jí)優(yōu)化。Relay, XLA, Torch-MLIR，ONXX 等基本都在這一級(jí)別。
張量程序表示(tensor program): 在這個(gè)級(jí)別我們需要對(duì)子圖進(jìn)行循環(huán)優(yōu)化，對(duì)于DSA支持還要包含張量化和內(nèi)存搬移的優(yōu)化。
算子庫和運(yùn)行環(huán)境(library and runtime): 算子庫本身依然是我們快速引入專家輸入優(yōu)化性能的方式。同時(shí)運(yùn)行環(huán)境快速支持?jǐn)?shù)據(jù)結(jié)構(gòu)運(yùn)行庫。
硬件專用指令 (hardware primitive) ：專用硬件和可編程深度學(xué)習(xí)加速器也引入了我們專用硬件張量指令的需求。
當(dāng)然的深度學(xué)習(xí)編譯生態(tài)基本上也是圍繞著各種抽象的實(shí)現(xiàn)展開。上圖是對(duì)于當(dāng)前生態(tài)的一個(gè)不完全總結(jié)。值得指出的是，算子庫和運(yùn)行環(huán)境等本身未必直接和編譯技術(shù)相關(guān)。但是因?yàn)槲覀兊哪繕?biāo)是新硬件部署和運(yùn)行加速，運(yùn)行環(huán)境抽象和硬件指令也是生態(tài)的重要一環(huán)。我們需要多層抽象本身基本是深度學(xué)習(xí)編譯和優(yōu)化領(lǐng)域大家達(dá)成的一種共識(shí)。但是為了真正地支持機(jī)器學(xué)習(xí)，光依賴于各個(gè)組件本身是遠(yuǎn)遠(yuǎn)不夠的。我們發(fā)現(xiàn)最大的問題是“如何設(shè)計(jì)各個(gè)層級(jí)的抽象，并且對(duì)它們進(jìn)行有效的整合”。目前的解決方案
如果我們仔細(xì)地研究目前的整個(gè)生態(tài)，包括深度學(xué)習(xí)框架，編譯框架(包括基于MLIR，ONNX或者是TVM的解決方案）。大家都遵循一種叫做多層漸進(jìn)優(yōu)化(Multi-stage lowering)的方式。這種構(gòu)建方式的大致思路是我們?cè)诿恳粚映橄笾胁捎靡粋€(gè)（有時(shí)多個(gè)）中間表示。我們會(huì)在每一個(gè)層級(jí)(dialect, abstraction）做一些內(nèi)部?jī)?yōu)化，然后把問題丟給下一個(gè)層級(jí)繼續(xù)進(jìn)行優(yōu)化。因?yàn)樯鷳B(tài)本身設(shè)計(jì)和其他一些原因，當(dāng)然的解決方案基本有以下特點(diǎn)：每一個(gè)層級(jí)抽象基本由一個(gè)比較獨(dú)立的團(tuán)體維護(hù)。層和層之間往往比較松耦合。最后解決方案本身往往是以一個(gè)黑盒工具的方式呈現(xiàn)給用戶。很多人的一個(gè)愿景是只要每一層的優(yōu)化做的好，把東西拼起來，我們就可以組合成一個(gè)滿足需求的解決方案。雖然這樣的做法的確可以達(dá)到一定的效果，但是我們也需要反問，multi-staging lowering真的可足以解決深度學(xué)習(xí)優(yōu)化的問題嗎?兩種隔閡
我們大約在三年前完成了基于multi-staging lowering的解決方案。當(dāng)我們把全棧解決方案搭建起來并且不斷實(shí)踐之后我們發(fā)現(xiàn)有兩種隔閡阻礙整個(gè)行業(yè)的發(fā)展。其中的第一種豎向隔閡阻隔了手工優(yōu)化的方案和自動(dòng)編譯優(yōu)化的方案。如果我們看當(dāng)前的深度學(xué)習(xí)運(yùn)行框架和編譯框架。我們可以發(fā)現(xiàn)有兩種流派：一類是以手工算子優(yōu)化為主的算子庫驅(qū)動(dòng)方案。這一類方案一般可以比較容易地讓更多的優(yōu)化專家加入，但是本身也會(huì)帶來比較多的工程開銷。另一類方案是以自動(dòng)優(yōu)化為主的編譯方案。編譯核心方案往往可以帶來更多的自動(dòng)的效果，但是有時(shí)候也比較難以引入領(lǐng)域知識(shí)。大部分當(dāng)前的框架基本都只為兩者中的其中一個(gè)設(shè)計(jì)。而我們往往發(fā)現(xiàn)實(shí)際比較好的解決方案其實(shí)同時(shí)需要機(jī)器學(xué)習(xí)工程師的輸入和自動(dòng)化。并且隨著領(lǐng)域的發(fā)展，最優(yōu)的解決方案也會(huì)發(fā)生變化。怎么樣打破這樣的豎向墻，讓手工優(yōu)化，機(jī)器學(xué)習(xí)優(yōu)化專家的知識(shí)和自動(dòng)優(yōu)化做有機(jī)整合，也是目前行業(yè)面臨的一個(gè)大的問題。
除了豎向高墻之外，第二類隔閡則一定程度上和multi-stage lowering的生態(tài)直接相關(guān)。因?yàn)槲覀兺鶗?huì)把不同層級(jí)的抽象分開來設(shè)計(jì)，雖然在抽象內(nèi)部可以做比較靈活的優(yōu)化。但是在一個(gè)抽象到另外一個(gè)抽象的轉(zhuǎn)換時(shí)候往往需要通過translator或者lowering批量轉(zhuǎn)換。這樣導(dǎo)致了我們很多困難都開始集中在一類抽象和另外一類抽象的邊界上，這樣導(dǎo)致了如果我們想要在邊界上面做一些分步的優(yōu)化（如把其中一部分子圖交給一類編譯邏輯，剩下的交給其他編譯如邏輯）我們就必須在邊界上面引入大量的工程。另外一個(gè)常見的現(xiàn)象是這類轉(zhuǎn)換往往是單向的，我們一般會(huì)在高階的抽象如計(jì)算圖上面做一些優(yōu)化，然后傳遞給張量計(jì)算層級(jí)。但是張量計(jì)算或者硬件層級(jí)的信息往往難以反饋給更高的層級(jí)。舉個(gè)例子，其實(shí)很多時(shí)候張量程序的優(yōu)化本身可以反過來指導(dǎo)計(jì)算圖層級(jí)的算子融合和數(shù)據(jù)排布，但是當(dāng)前的單向架構(gòu)比較難自然地利用這一類反饋。總結(jié)一下我們的經(jīng)驗(yàn)。深度學(xué)習(xí)編譯和優(yōu)化本身不是一個(gè)一個(gè)層級(jí)可以全部完成優(yōu)化的問題。解決相關(guān)問題需要各個(gè)層級(jí)抽象之間的聯(lián)動(dòng)。隨著TVM和MLIR一類基礎(chǔ)架構(gòu)的出現(xiàn)，我們其實(shí)已經(jīng)可以比較容易地搭建出某一個(gè)層級(jí)的抽象或者dialect并且讓它們之間通過multi-stage lowering的方式從高到地級(jí)抽象進(jìn)行逐層變換和轉(zhuǎn)換。但是現(xiàn)在的困難點(diǎn)往往出現(xiàn)在抽象的轉(zhuǎn)換邊界上。不論是在邊界引入更多的可模塊化整合變換，或者是進(jìn)行嘗試反饋迭代，multi-stage lowering本身是遠(yuǎn)遠(yuǎn)不夠的。不僅如此，隨著抽象層級(jí)的增加，如果希望加入自定義算子，我們往往需要針對(duì)每個(gè)抽象層級(jí)進(jìn)行進(jìn)一步的架構(gòu)，其中的開銷變得更加龐大。因?yàn)楦鱾€(gè)抽象之間的豎向和橫向隔閡。不論我們?nèi)绾巫龊靡粚映橄髢?nèi)部本身，我們依然難以做好端到端的整體優(yōu)化。需要注意的是，這些隔閡和問題的存在和基礎(chǔ)架構(gòu)的選擇無關(guān)，不論是基于MLIR，ONNX或者是TVM的方案，一旦采用了multi-stage lowering都會(huì)不可避免地面對(duì)這個(gè)問題。相信在領(lǐng)域里面努力的小伙伴不管采取什么基礎(chǔ)架構(gòu)，在把方案打通之后或多或少都會(huì)碰到這個(gè)本質(zhì)的問題。未來在哪里：從箭頭到圈
這一節(jié)我們會(huì)介紹對(duì)于這一個(gè)問題經(jīng)過兩年多總結(jié)的思考。這也是我們演化到新一代深度學(xué)習(xí)編譯系統(tǒng)的核心技術(shù)路線，我們把這一路線叫做TVM Unity。我們可以注意到，幾乎所有的難點(diǎn)都由邊界隔閡產(chǎn)生，因此我們需要解決的重點(diǎn)是抓住關(guān)鍵點(diǎn)，消除邊界隔閡。我們的目標(biāo)是把一個(gè)單向箭頭的multi-stage lowering方案，演化成一個(gè)可以讓各個(gè)抽象之間有機(jī)交互的一個(gè)圈。
整個(gè)技術(shù)路線總結(jié)下來包含三大關(guān)鍵點(diǎn):Unify: 統(tǒng)一多層抽象
為了打破抽象直接隔閡，我們需要完成的第一步是統(tǒng)一抽象。這當(dāng)然并不意味著我們需要設(shè)計(jì)一個(gè)唯一的層級(jí)來解決所有問題 – 架構(gòu)上面稍有不慎，我們可能會(huì)設(shè)計(jì)出一個(gè)整合了所有抽象短板的復(fù)雜表示。在這里我們依然需要承認(rèn)每一類抽象的重要性，但是我們需要在不同類別的抽象之間進(jìn)行協(xié)同設(shè)計(jì)，并且可以讓每一個(gè)層級(jí)和其它層級(jí)進(jìn)行相互交互。
具體到TVM而言，TVM主要的重點(diǎn)放在四個(gè)抽象上。AutoTensorization用來解決硬件指令生命和張量程序?qū)?，TVM FFI（PackedFunc）機(jī)制使得我們可以靈活地引入任意的算子庫和運(yùn)行庫函數(shù)并且在各個(gè)編譯模塊和自定義模塊里面相互調(diào)用。TensorIR負(fù)責(zé)張量級(jí)別程序和硬件張量指令的整合。Relax (Relax Next) 會(huì)引入relay的進(jìn)一步迭代，直接引入first class symbolic shape的支持。但是就和一開始提到的一樣，這里的關(guān)鍵點(diǎn)并不只在各個(gè)抽象層級(jí)本身，而是抽象之間的相互交互和聯(lián)合優(yōu)化。
以上這個(gè)程序樣例就展示了我們統(tǒng)一抽象的設(shè)計(jì)目標(biāo)。這是一個(gè)通過TVMScript表示的IRModule。MyIRModule是包含了兩種函數(shù)。其中tir_mm是一個(gè)TensorIR級(jí)別的函數(shù)。新一代的TensorIR的設(shè)計(jì)目標(biāo)是實(shí)現(xiàn) 張量程序和硬件專用指令之間的通過自動(dòng)張量化的聯(lián)動(dòng)。再看relax_func這個(gè)函數(shù)。其中有幾個(gè)關(guān)鍵點(diǎn)。R.call_dps((n, m), tir_mm, [x, w])是在計(jì)算圖級(jí)別直接對(duì)于TensorIR函數(shù)tir_mm的直接調(diào)用。并且我們通過特殊的調(diào)用形式使得張量程序依然以和圖算子一樣的方式出現(xiàn)在計(jì)算圖中。支持這一特性意味著計(jì)算圖需要和張量程序?qū)蛹?jí)表示進(jìn)行聯(lián)合設(shè)計(jì)，使得計(jì)算圖的優(yōu)化可以利用張量程序?qū)蛹?jí)的信息。最后 R.call_packed("custom_inplace_update", gv0) 允許計(jì)算圖直接和TVM FFI函數(shù)進(jìn)行交互。多層抽象之間的相互統(tǒng)一整合雖然帶來了更多的一些設(shè)計(jì)考慮，但是也帶來了解決隔閡之后的不少優(yōu)勢(shì)。舉個(gè)例子，假設(shè)我們現(xiàn)在有一個(gè)快速的算子優(yōu)化思路，在傳統(tǒng)的編譯流程中常見的做法是引入在每個(gè)層級(jí)引入新的改寫規(guī)則。但是在統(tǒng)一抽象下，我們可以直接引入一個(gè)pass把一個(gè)局部算子先改寫成call_packed調(diào)用一個(gè)手寫的外部算子。在確認(rèn)性能合理的情況下再考慮如何把一個(gè)局部改寫變成更加自動(dòng)化的方案。同樣我們也可以把手工算子和自動(dòng)代碼生成的方案更加有機(jī)地整合在一起。最后，因?yàn)閷?duì)于TensorIR程序的調(diào)用直接被表示再圖程序中，我們可以直接把對(duì)于TensorIR的改寫優(yōu)化變成對(duì)于對(duì)應(yīng)計(jì)算圖的調(diào)用改寫，使得張量級(jí)別變換的信息可以反饋到圖級(jí)別的變換中去。最后，值得注意的是樣例中關(guān)于動(dòng)態(tài)shape的支持采用了symbolic shape的方案。樣例中的(n, m)和(n * m)中的n和m貫穿整個(gè)程序。使得我們可以在編譯優(yōu)化時(shí)利用這些更多的信息來做更多動(dòng)態(tài)相關(guān)優(yōu)化。并且關(guān)于symbolic表達(dá)式的支持也完美地和TensorIR層級(jí)的symbolic shape統(tǒng)一在了一起。這一特性也很大地體現(xiàn)和協(xié)同抽象設(shè)計(jì)的重要性。Interact: 交互開放迭代
除了抽象的整合之外，一個(gè)非常重要的課題是如何讓不同的人之間進(jìn)行相互協(xié)作。深度學(xué)習(xí)編譯優(yōu)化是一個(gè)涉及到各種各樣工程師人群的領(lǐng)域。算法專家希望開發(fā)新的模型和自定義算子，機(jī)器學(xué)習(xí)系統(tǒng)工程師需要開發(fā)系統(tǒng)優(yōu)化，硬件廠商需要迭代自己的硬件。每一類人都會(huì)變對(duì)不同層級(jí)的抽象。深度學(xué)習(xí)領(lǐng)域的繁榮很大程度地歸功于深度學(xué)習(xí)框架的開放生態(tài)。任何人都可以以python的方式編寫我們的模型并且把它們和其人寫的模塊進(jìn)行整合。傳統(tǒng)的編譯器領(lǐng)域往往會(huì)一個(gè)更加封閉的方式呈現(xiàn)出來。在一個(gè)multi-stage lowering的架構(gòu)下把框架搭好，然后提供一個(gè)命令行接口給用戶。雖然一個(gè)美好的愿景是我們可以把每一個(gè)層級(jí)的抽象做好最好再，然后再把火箭的零件全部拼接起來。但是這樣的愿望往往是不現(xiàn)實(shí)的，實(shí)際的情況是每一層都還是會(huì)有自己的問題。如何通過系統(tǒng)工程的方式做好整合互補(bǔ)，并且可以讓不同的人群可以在一起相互快速合作迭代出新的解決方案才是我們需要考慮的問題。
而這個(gè)時(shí)候允許不同的人之間進(jìn)行協(xié)作，交互開發(fā)和迭代就是一個(gè)我們需要首要考慮的話題。在這個(gè)層面上，我們遵循以下幾點(diǎn)原則: python-first，通過TVMscript和直接多語言整合的架構(gòu)讓大家都可以通過python API相互協(xié)作。開放開發(fā)不論是哪個(gè)抽象層級(jí)之間都可以整合交流。協(xié)作迭代，讓大家可以一起合作達(dá)到更好的效果。舉個(gè)例子，一個(gè)機(jī)器學(xué)習(xí)專家可以利用計(jì)算表達(dá)式來寫一個(gè)自定義算子，但這個(gè)自定義算子可以被系統(tǒng)專家寫的自動(dòng)轉(zhuǎn)換規(guī)則優(yōu)化，而其中自動(dòng)轉(zhuǎn)換規(guī)則本身又利用的硬件廠商所提供的張量指令特性。當(dāng)然這個(gè)整個(gè)鏈路的每一環(huán)都可以在一起快速聯(lián)動(dòng)的時(shí)候，我們就可以快速地根據(jù)需求迭代出想要的解決方案。Automate: 自動(dòng)優(yōu)化整合
自動(dòng)優(yōu)化始終是深度學(xué)習(xí)編譯器血液里面的東西，也是實(shí)現(xiàn)Unity的一個(gè)關(guān)鍵環(huán)節(jié)。
很多傳統(tǒng)的自動(dòng)優(yōu)化方案是排他的，也就是只有整個(gè)優(yōu)化方案全部采用對(duì)應(yīng)的模型的做法（如一些polyhedral model）才可以解決好，一旦嘗試引入領(lǐng)域知識(shí)，或者想要的優(yōu)化跳出了原有的范疇，自動(dòng)化就難以發(fā)揮優(yōu)勢(shì)。我們需要改變這一觀點(diǎn)，重點(diǎn)思考如何在自動(dòng)優(yōu)化的過程中有效地整合領(lǐng)域?qū)＜抑R(shí)。使得我們真正可以把自動(dòng)化和高手之間的力量整合在一起。社區(qū)基于TensorIR的MetaSchedule正是往這個(gè)方向前進(jìn)的重要一步。前面的幾節(jié)我們介紹了TVM Unity的三個(gè)關(guān)鍵技術(shù)點(diǎn)。當(dāng)然unity本身設(shè)計(jì)的目標(biāo)本身并不是解決所有問題。我們非常清楚只有當(dāng)把這個(gè)圈和整個(gè)機(jī)器學(xué)習(xí)和硬件生態(tài)一起整合的時(shí)候我們才可以發(fā)揮最大的效率。而抽象的整合使得我們可以更加容易地提供更多種整合方式。我們可以通過子圖整合把TVM整合入已有的框架中，也可以把常見的硬件后端整合入TVM本身。統(tǒng)一的抽象可以使得這樣的整合以相對(duì)統(tǒng)一的方式發(fā)生在各個(gè)層級(jí)。
TVM FFI 等靈活的接口也讓我們可以和整個(gè)生態(tài)做靈活的對(duì)接。并且把不同的后端生態(tài)通過自動(dòng)化做更加有效的整合。總結(jié)和未來展望
這篇文章總結(jié)了我們對(duì)于深度學(xué)習(xí)編譯領(lǐng)域在過去兩年的思考和未來的展望。對(duì)于新一代架構(gòu)的推動(dòng)一直是我們核心關(guān)注的主題，這里提到的各個(gè)特性也都已經(jīng)重構(gòu)完成或者在進(jìn)行中。TVM FFI在去年已經(jīng)逐漸成熟，TensorIR本身剛被合并到主干，后續(xù)的metaschedule也會(huì)陸續(xù)進(jìn)入主干。Relax本身也在社區(qū)進(jìn)行開放開發(fā)。TVM Unity作為社區(qū)的核心主題也會(huì)是接下來一年的重點(diǎn)，并且在開發(fā)過程中的一些元素已經(jīng)可以提供不少好處。當(dāng)社區(qū)逐漸向這一技術(shù)方向靠攏對(duì)接，整合的優(yōu)勢(shì)也會(huì)越來越明顯。過去一年中常常看到一些關(guān)于深度學(xué)習(xí)編譯基礎(chǔ)架構(gòu)的討論。其實(shí)從核心上面來看，基礎(chǔ)架構(gòu)本身當(dāng)然是重要的一環(huán)（scala之于spark一樣），不過不論MLIR, TVM或者其他編譯框架基礎(chǔ)架構(gòu)本身也都在相互學(xué)習(xí)中趨近成熟。真正的瓶頸還是在于抽象的設(shè)計(jì)和更進(jìn)一步真正地進(jìn)行協(xié)同整合。當(dāng)然其實(shí)三個(gè)技術(shù)關(guān)鍵點(diǎn)的選擇也的確影響了我們對(duì)于基礎(chǔ)架構(gòu)的思考。比如TVM FFI作為比較重要的基礎(chǔ)架構(gòu)支撐了交互和python-first的特性。不論基礎(chǔ)架構(gòu)，我們面對(duì)的真正問題是如何解決這些關(guān)鍵的隔閡問題，multi-stage lowering本身的缺陷導(dǎo)致了現(xiàn)有的方案必須有所突破才能演化到新一代的深度學(xué)習(xí)編譯系統(tǒng)。本文探討了這一演化的方向和具體的技術(shù)路線。希望可以對(duì)整個(gè)領(lǐng)域有所啟發(fā)。在新一代的架構(gòu)前進(jìn)的道路上，我們不再是一個(gè)人在努力。很多的關(guān)鍵組件設(shè)計(jì)都是由各個(gè)同學(xué)一起合作完成。我們也歡迎更多的同學(xué)一起加入到社區(qū)中來，一起推動(dòng)實(shí)現(xiàn)新一代深度學(xué)習(xí)編譯系統(tǒng)的共同建設(shè)中來。原文鏈接：https://zhuanlan.zhihu.com/p/446935289?utm_source=wechat_session&utm_medium=social&s_r=0本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。
