一導讀

隨著深度學習的不斷發展，AI模型結構在快速演化，底層計算硬體技術更是層出不窮，對於廣大開發者來說不僅要考慮如何在複雜多變的場景下有效的將算力發揮出來，還要應對計算框架的持續迭代。深度編譯器就成了應對以上問題廣受關注的技術方向，讓使用者僅需專注於上層模型開發，降低手工最佳化效能的人力開發成本，進一步壓榨硬體效能空間。阿里雲機器學習PAI開源了業內較早投入實際業務應用的動態shape深度學習編譯器 BladeDISC，本文將詳解 BladeDISC的設計原理和應用。

二 BladeDISC是什麼

BladeDISC是阿里最新開源的基於MLIR的動態shape深度學習編譯器。

1 主要特性

多款前端框架支援：TensorFlow，PyTorch
多後端硬體支援：CUDA，ROCM，x86

完備支援動態shape語義編譯
支援推理及訓練

輕量化API，對使用者通用透明
支援外掛模式嵌入宿主框架執行，以及獨立部署模式

三深度學習編譯器的背景

近幾年來，深度學習編譯器作為一個較新的技術方向異常地活躍，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，到後來呼聲很高的MLIR以及其不同領域的延伸專案IREE、mlir-hlo等等。能夠看到不同的公司、社群在這個領域進行著大量的探索和推進。

1 AI浪潮與晶片浪潮共同催生——從萌芽之初到蓬勃發展

深度學習編譯器近年來之所以能夠受到持續的關注，主要來自於幾個方面的原因：

框架效能最佳化在模型泛化性方面的需求

深度學習還在日新月異的發展之中，創新的應用領域不斷湧現，複雜多變的場景下如何有效的將硬體的算力充分發揮出來成為了整個AI應用鏈路中非常重要的一環。在早期，神經網路部署的側重點在於框架和運算元庫，這部分職責很大程度上由深度學習框架、硬體廠商提供的運算元庫、以及業務團隊的手工最佳化工作來承擔。

上圖將近年的深度學習框架粗略分為三代。一個趨勢是在上層的使用者API層面，這些框架在變得越來越靈活，靈活性變強的同時也為底層效能問題提出了更大的挑戰。初代深度學習框架類似 Caffe 用 sequence of layer 的方式描述神經網路結構，第二代類似 TensorFlow 用更細粒度的 graph of operators 描述計算圖，到第三代類似 PyTorch，TensorFlow Eager Mode 的動態圖。我們可以看到框架越來越靈活，描述能力越來越強，帶來的問題是最佳化底層效能變得越來越困難。業務團隊也經常需要補充完成所需要的手工最佳化，這些工作依賴於特定業務和對底層硬體的理解，耗費人力且難以泛化。而深度學習編譯器則是結合編譯時圖層的最佳化以及自動或者半自動的程式碼生成，將手工最佳化的原理做泛化性的沉澱，以替代純手工最佳化帶來的各種問題，去解決深度學習框架的靈活性和效能之間的矛盾。

AI框架在硬體泛化性方面的需求

表面上看近些年AI發展的有目共睹、方興未艾，而後臺的硬體算力數十年的發展才是催化AI繁榮的核心動力。隨著電晶體微縮面臨的各種物理挑戰越來越大，晶片算力的增加越來越難，摩爾定律面臨失效，創新體系結構的各種DSA晶片迎來了一波熱潮，傳統的x86、ARM等都在不同的領域強化著自己的競爭力。硬體的百花齊放也為AI框架的發展帶來了新的挑戰。

而硬體的創新是一個問題，如何能把硬體的算力在真實的業務場景中發揮出來則是另外一個問題。新的AI硬體廠商不得不面臨除了要在硬體上創新，還要在軟體棧上做重度人力投入的問題。向下如何相容硬體，成為了如今深度學習框架的核心難點之一，而相容硬體這件事，則需要由編譯器來解決。

AI系統平臺對前端AI框架泛化性方面的需求

當今主流的深度學習框架包括Tensoflow、Pytorch、Keras、JAX等等，幾個框架都有各自的優缺點，在上層對使用者的介面方面風格各異，卻同樣面臨硬體適配以及充分發揮硬體算力的問題。不同的團隊常根據自己的建模場景和使用習慣來選擇不同的框架，而云廠商或者平臺方的效能最佳化工具和硬體適配方案則需要同時考慮不同的前端框架，甚至未來框架演進的需求。Google利用XLA同時支援TensorFlow和JAX，同時其它開源社群還演進出了Torch_XLA，Torch-MLIR這樣的接入方案，這些接入方案目前雖然在易用性和成熟程度等方面還存在一些問題，卻反應出AI系統層的工作對前端AI框架泛化性方面的需求和技術趨勢。

2 什麼是深度學習編譯器

傳統編譯器是以高層語言作為輸入，避免使用者直接去寫機器碼，而用相對靈活高效的語言來工作，並在編譯的過程中引入最佳化來解決高層語言引入的效能問題，平衡開發效率與效能之間的矛盾。而深度學習編譯器的作用相仿，其輸入是比較靈活的，具備較高抽象度的計算圖描述，輸出包括CPU、GPU及其他異構硬體平臺上的底層機器碼及執行引擎。

傳統編譯器的使命之一是減輕程式設計者的壓力。作為編譯器的輸入的高階語言往往更多地是描述一個邏輯，為了便利程式設計者，高階語言的描述會更加抽象和靈活，至於這個邏輯在機器上是否能夠高效的執行，往往是考驗編譯器的一個重要指標。深度學習作為一個近幾年發展異常快速的應用領域，它的效能最佳化至關重要，並且同樣存在高層描述的靈活性和抽象性與底層計算效能之間的矛盾，因此專門針對深度學習的編譯器出現了。而傳統編譯器的另外一個重要使命是，需要保證程式設計者輸入的高層語言能夠執行在不同體系結構和指令集的硬體計算單元上，這一點也同樣反應在深度學習編譯器上。面對一個新的硬體裝置，人工的話不太可能有精力對那麼多目標硬體重新手寫一個框架所需要的全部運算元實現，深度學習編譯器提供中間層的IR，將頂層框架的模型流圖轉化成中間層表示IR，在中間層IR上進行通用的圖層最佳化，而在後端將最佳化後的IR通用性的生成各個目標平臺的機器碼。

深度學習編譯器的目標是針對AI計算任務，以通用編譯器的方式完成效能最佳化和硬體適配。讓使用者可以專注於上層模型開發，降低使用者手工最佳化效能的人力開發成本，進一步壓榨硬體效能空間。

3 距離大規模應用面臨的瓶頸問題

深度學習編譯器發展到今天，雖然在目標和技術架構方面與傳統編譯器有頗多相似之處，且在技術方向上表現出了良好的潛力，然而目前的實際應用範圍卻仍然距離傳統編譯器有一定的差距，主要難點包括：

易用性

深度學習編譯器的初衷是希望簡化手工最佳化效能和適配硬體的人力成本。然而在現階段，大規模部署應用深度學習編譯器的挑戰仍然較大，能夠將編譯器用好的門檻較高，造成這個現象的主要原因包括：

與前端框架對接的問題。不同框架對深度學習任務的抽象描述和API介面各有不同，語義和機制上有各自的特點，且作為編譯器輸入的前端框架的運算元型別數量呈開放性狀態。如何在不保證所有運算元被完整支援的情況下透明化的支援使用者的計算圖描述，是深度學習編譯器能夠易於為使用者所廣泛使用的重要因素之一。

動態shape問題和動態計算圖問題。現階段主流的深度學習編譯器主要針對特定的靜態shape輸入完成編譯，此外對包含control flow語義的動態計算圖只能提供有限的支援或者完全不能夠支援。而AI的應用場景卻恰恰存在大量這一類的任務需求。這時只能人工將計算圖改寫為靜態或者半靜態的計算圖，或者想辦法將適合編譯器的部分子圖提取出來交給編譯器。這無疑加重了應用深度學習編譯器時的工程負擔。更嚴重的問題是，很多工型別並不能透過人工的改寫來靜態化，這導致這些情況下編譯器完全無法實際應用。

編譯開銷問題。作為效能最佳化工具的深度學習編譯器只有在其編譯開銷對比帶來的效能收益有足夠優勢的情況下才真正具有實用價值。部分應用場景下對於編譯開銷的要求較高，例如普通規模的需要幾天時間完成訓練任務有可能無法接受幾個小時的編譯開銷。對於應用工程師而言，使用編譯器的情況下不能快速的完成模型的除錯，也增加了開發和部署的難度和負擔。

對使用者透明性問題。部分AI編譯器並非完全自動的編譯工具，其效能表現比較依賴於使用者提供的高層抽象的實現模版。主要是為運算元開發工程師提供效率工具，降低使用者人工調優各種運算元實現的人力成本。但這也對使用者的運算元開發經驗和對硬體體系結構的熟悉程度提出了比較高的要求。此外，對於新硬體的軟體開發者來說，現有的抽象卻又常常無法足夠描述創新的硬體體系結構上所需要的運算元實現。需要對編譯器架構足夠熟悉的情況下對其進行二次開發甚至架構上的重構，門檻及開發負擔仍然很高。

魯棒性

目前主流的幾個AI編譯器專案大多數都還處於偏實驗性質的產品，但產品的成熟度距離工業級應用有較大的差距。這裡的魯棒性包含是否能夠順利完成輸入計算圖的編譯，計算結果的正確性，以及效能上避免coner case下的極端badcase等。

效能問題

編譯器的最佳化本質上是將人工的最佳化方法，或者人力不易探究到的最佳化方法透過泛化性的沉澱和抽象，以有限的編譯開銷來替代手工最佳化的人力成本。然而如何沉澱和抽象的方法學是整個鏈路內最為本質也是最難的問題。深度學習編譯器只有在效能上真正能夠代替或者超過人工最佳化，或者真正能夠起到大幅度降低人力成本作用的情況下才能真正發揮它的價值。

然而達到這一目標卻並非易事，深度學習任務大都是tensor級別的計算，對於並行任務的拆分方式有很高的要求，但如何將手工的最佳化泛化性的沉澱在編譯器技術內，避免編譯開銷爆炸以及分層之後不同層級之間最佳化的聯動，仍然有更多的未知需要去探索和挖掘。這也成為以MLIR框架為代表的下一代深度學習編譯器需要著重思考和解決的問題。

四 BladeDISC的主要技術特點

專案最早啟動的初衷是為了解決XLA和TVM當前版本的靜態shape限制，內部命名為 DISC (DynamIc Shape Compiler)，希望打造一款能夠在實際業務中使用的完備支援動態shape語義的深度學習編譯器。

從團隊四年前啟動深度學習編譯器方向的工作以來，動態shape問題一直是阻礙實際業務落地的嚴重問題之一。彼時，包括XLA在內的主流深度學習框架，都是基於靜態shape語義的編譯器框架。典型的方案是需要使用者指定輸入的shape，或是由編譯器在執行時捕捉待編譯子圖的實際輸入shape組合，並且為每一個輸入shape組合生成一份編譯結果。

靜態shape編譯器的優勢顯而易見，編譯期完全已知靜態shape資訊的情況下，Compiler可以作出更好的最佳化決策並得到更好的CodeGen效能，同時也能夠得到更好的視訊記憶體/記憶體最佳化plan和排程執行plan。然而，其缺點也十分明顯，具體包括：

大幅增加編譯開銷。引入離線編譯預熱過程，大幅增加推理任務部署過程複雜性；訓練迭代速度不穩定甚至整體訓練時間負最佳化。

部分業務場景shape變化範圍趨於無窮的，導致編譯快取永遠無法收斂，方案不可用。

記憶體視訊記憶體佔用的增加。編譯快取額外佔用的記憶體視訊記憶體，經常導致實際部署環境下的記憶體/視訊記憶體OOM，直接阻礙業務的實際落地。

人工padding為靜態shape等緩解性方案對使用者不友好，大幅降低應用的通用性和透明性，影響迭代效率。

在2020年夏天，DISC完成了僅支援TensorFlow前端以及Nvidia GPU後端的初版，並且正式在阿里內部上線投入實際應用。最早在幾個受困於動態shape問題已久的業務場景上投入使用，並且得到了預期中的效果。即在一次編譯且不需要使用者對計算圖做特殊處理的情況下，完備支援動態shape語義，且效能幾乎與靜態shape編譯器持平。對比TensorRT等基於手工運算元庫為主的最佳化框架，DISC基於編譯器自動codegen的技術架構在經常為非標準開源模型的實際業務上獲得了明顯的效能和易用性優勢。

從2020年第二季度開始至今，DISC持續投入研發力量，針對前文提到的從雲端平臺方視角看到的深度學習編譯器距離大規模部署和應用的幾個瓶頸問題，在效能、運算元覆蓋率和魯棒性、CPU及新硬體支援、前端框架支援等方面逐漸完善。目前在場景覆蓋能力和效能等方面，已經逐漸替換掉團隊過往基於XLA和TVM等靜態shape框架上的工作，成為PAI-Blade支援阿里內部及阿里雲外部業務的主要最佳化手段。2021年後，DISC在CPU及GPGPU體系結構的後端硬體上的效能有了顯著的提升，同時在新硬體的支援上面投入了更多的技術力量。2021年底，為了吸引更多的技術交流和合作共建需要，以及更大範圍的使用者反饋，正式更名為BladeDISC並完成了初版開源。

五關鍵技術

BladeDISC的整體架構，及其在阿里雲相關產品中的上下文關係如下圖所示：

1 MLIR基礎架構

MLIR是由Google在2019年發起的專案，MLIR 的核心是一套靈活的多層IR基礎設施和編譯器實用工具庫，深受 LLVM 的影響，並重用其許多優秀理念。這裡我們選擇基於MLIR的主要原因包括其比較豐富的基礎設施支援，方便擴充套件的模組化設計架構以及MLIR較強的膠水能力。

2 動態shape編譯

上圖為BladeDISC的主體Pass Pipeline設計。對比目前主流的深度學習編譯器專案，主要技術特點如下：

圖層IR設計

BladeDISC選擇基於HLO作為核心圖層IR來接入不同的前端框架，但是HLO是原本為XLA設計的純靜態shape語義的IR。靜態場景下，HLO IR中的shape表達會被靜態化，所有的shape計算會被固化為編譯時常量保留在編譯結果中；而在動態shape場景下，IR本身需要有足夠的能力表達shape計算和動態shape資訊的傳遞。BladeDISC從專案建立開始一直與MHLO社群保持緊密的合作，在XLA的HLO IR基礎上，擴充套件了一套具有完備動態shape表達能力的IR，並增加了相應的基礎設施以及前端框架的運算元轉換邏輯。這部分實現目前已經完整upstream至MHLO社群，確保後續其它MHLO相關專案中IR的一致性。

執行時Shape計算、儲存管理和Kernel排程

動態shape編譯的主要挑戰來自於需要在靜態的編譯過程中能夠處理動態的計算圖語義。為完備支援動態shape，編譯結果需要能夠在執行時做即時的shape推導計算，不僅要為資料計算，同時也需要為shape計算做程式碼生成。計算後的shape資訊用於做記憶體/視訊記憶體管理，以及kernel排程時的引數選擇等等。BladeDISC的pass pipeline的設計充分考慮了上述動態shape語義支援的需求，採用了host-device聯合codegen的方案。以GPU Backend為例，包括shape計算、記憶體/視訊記憶體申請釋放、硬體管理、kernel launch執行時流程全部為自動程式碼生成，以期得到完備的動態shape端到端支援方案和更為極致的整體效能。

動態shape下的效能問題

在shape未知或者部分未知的情況下，深度學習編譯器在效能上面臨的挑戰被進一步放大。在大多數主流硬體backend上，BladeDISC採用區分計算密集型部分和訪存密集型部分的策略，以期在效能與複雜性和編譯開銷之間獲取更好的平衡。

對於計算密集型部分，不同的shape要求更加精細的schedule實現來獲得更好的效能，pass pipeline在設計上的主要考慮是需要支援在執行時根據不同的具體shape選擇合適的運算元庫實現，以及處理動態shape語義下的layout問題。

而訪存密集型部分的自動運算元融合作為深度學習編譯器主要的效能收益來源之一，同樣面臨shape未知情況下在效能上的挑戰。許多靜態shape語義下比較確定性的問題，例如指令層的向量化，codegen模版選擇，是否需要implicit broadcast等等在動態shape場景下都會面臨更大的複雜性。針對這些方面的問題，BladeDISC選擇將部分的最佳化決策從編譯時下沉到執行時。即在編譯期根據一定的規則生成多個版本的kernel實現，在執行時根據實際shape自動選擇最優的實現。這一機制被稱作speculation，在BladeDISC內基於host-device的聯合程式碼生成來實現。此外，在編譯期沒有具體shape數值的情況下，會很容易在各個層級丟失掉大量的最佳化機會，從圖層的線性代數簡化、fusion決策到指令層級的CSE、常數摺疊等。BladeDISC在IR及pass pipeline的設計過程中著重設計了shape constraint在IR中的抽象和在pass pipeline中的使用，例如編譯期未知的不同dimension size之間的約束關係等。在最佳化整體效能方面起到了比較明顯的作用，保證能夠足夠接近甚至超過靜態shape編譯器的效能結果。

大顆粒度運算元融合

團隊在開啟BladeDISC專案之前，曾經基於靜態shape編譯器在大顆粒度運算元融合及自動程式碼生成方面有過若干探索[3][4]，其基本思想可以概括為藉助於GPU硬體中低訪存開銷的shared memory或CPU中低訪存開銷的Memory Cache，將不同schedule的計算子圖縫合進同一個kernel內，實現多個parallel loop複合，這種codegen方法稱之為fusion-stitching。這種訪存密集型子圖的自動程式碼生成打破了常規的loop fusion，input/output fusion對fusion顆粒度的限制。在保證程式碼生成質量的同時，大幅增加fusion顆粒度，同時避免複雜性及編譯開銷爆炸。且整個過程完全對使用者透明，無需人工指定schedule描述。

在動態shape語義下實現fusion-stitching對比靜態shape語義下同樣需要處理更大的複雜性，動態shape語義下的shape constraint抽象一定程度上簡化了這一複雜性，使整體效能進一步接近甚至超過手工運算元實現。

3 多前端框架支援

AICompiler框架在設計時也包含了擴充套件支援不同前端框架的考慮。PyTorch側透過實現一個輕量的Converter將TorchScript轉換為DHLO IR實現了對PyTorch推理作業的覆蓋。MLIR相對完備的IR基礎設施也為Converter的實現提供了便利。BladeDISC包含Compiler以及適配不同前端框架的Bridge側兩部分。其中Bridge進一步分為宿主框架內的圖層pass以及執行時Op兩部分，以外掛的方式接入宿主框架。這種工作方式使BladeDISC可以透明化的支援前端計算圖，可以適配使用者各種版本的宿主框架。

4 執行時環境適配

為將編譯的結果能夠配合TensorFlow/PyTorch等宿主在各自的執行環境中執行起來，以及管理執行時IR層不易表達的狀態資訊等等，我們為不同的執行時環境實現了一套統一的Compiler架構，並引入了執行時抽象層，即RAL(Runtime Abstraction Layer)層。

RAL實現了多種執行環境的適配支援，使用者可以根據需要進行選擇，具體包括：

全圖編譯，獨立執行。當整個計算圖都支援編譯時，RAL提供了一套簡易的runtime以及在此之上RAL Driver的實現，使得compiler編譯出來結果可以脫離框架直接執行，減少框架overhead。

TF中子圖編譯執行。

Pytorch中子圖編譯執行。

以上環境中在諸如資源管理，API語義等上存在差異，RAL透過抽象出一套最小集合的API ，並清晰的定義出它們的語義，將編譯器與執行時隔離開來，來達到在不同的環境中都能夠執行編譯出來的結果的目的。此外RAL層實現了無狀態編譯，解決了計算圖的編譯之後，編譯的結果可能被多次執行時的狀態資訊處理問題。一方面簡化了程式碼生成的複雜度，另一方面也更容易支援多執行緒併發執行（比如推理）的場景，同時在錯誤處理，回滾方面也更加容易支援。

六應用場景

BladeDISC的典型應用場景可以不太嚴格的分為兩類：其一是在主流的硬體平臺上（包括Nvidia GPU，x86 CPU等）上作為通用、透明的效能最佳化工具，降低使用者部署AI作業的人力負擔，提高模型迭代效率；另一個重要的應用場景是幫助新硬體做AI場景的適配和接入支援。

目前BladeDISC已經廣泛應用在阿里內部和阿里雲上外部使用者的多個不同應用場景下，覆蓋模型型別涉及NLP、機器翻譯、語音類ASR、語音TTS、影像檢測、識別、AI for science等等多種典型AI應用；覆蓋行業包括網際網路、電商、自動駕駛、安全行業、線上娛樂、醫療和生物等等。

在推理場景下，BladeDISC與TensorRT等廠商提供的推理最佳化工具形成良好的技術互補，其主要差異性優勢包括：

應對動態shape業務完備的動態shape語義支援

基於compiler based的技術路徑的模型泛化性在非標準模型上的效能優勢

更為靈活的部署模式選擇，以外掛形式支援前端框架的透明性優勢

下圖為Nvidia T4硬體上幾個真實的業務例子的效能收益數字：

在新硬體支援方面，目前普遍的情況是除了積累比較深厚的Nvidia等頭部廠商之外，包括ROCM等其它GPGPU硬體普遍存在的情況是硬體的指標已經具備相當的競爭力，但廠商受制於AI軟體棧上的積累相對較少，普遍存在硬體算力無法發揮出來導致硬體落地應用困難的問題。如前文所述，基於編譯器的技術路徑下天然對於硬體的後端具備一定的泛化能力，且與硬體廠商的技術儲備形成比較強的互補。BladeDISC目前在GPGPU和通用CPU體系結構上的儲備相對比較成熟。以GPGPU為例，在Nvidia GPU上的絕大部分技術棧可以遷移至海光DCU和AMD GPU等體系結構相近的硬體上。BladeDISC較強的硬體泛化能力配合硬體本身較強的通用性，很好的解決了新硬體適配的效能和可用性問題。

下圖為海光DCU上幾個真實業務例子上的效能數字：

某識別類模型	推理	不同batchsize下 2.21X ～ 2.31X
某檢測類模型A	推理	不同batchsize下 1.73X ～ 2.1X
某檢測類模型B	推理	不同batchsize下 1.04X ～ 1.59X
某分子動力學模型	訓練	2.0X

七開源生態——構想和未來

我們決定建設開源生態主要有如下的考慮：

BladeDISC發源於阿里雲計算平臺團隊的業務需求，在開發過程中與MLIR/MHLO/IREE等社群同行之間的討論和交流給了我們很好的輸入和借鑑。在我們自身隨著業務需求的迭代逐漸完善的同時，也希望能夠開源給社群，在目前整個AI編譯器領域實驗性專案居多，偏實用性強的產品偏少，且不同技術棧之間的工作相對碎片化的情況下，希望能夠將自身的經驗和理解也同樣回饋給社群，希望和深度學習編譯器的開發者和AI System的從業者之間有更多更好的交流和共建，為這個行業貢獻我們的技術力量。

我們希望能夠藉助開源的工作，收到更多真實業務場景下的使用者反饋，以幫助我們持續完善和迭代，併為後續的工作投入方向提供輸入。

後續我們計劃以兩個月為單位定期釋出Release版本。BladeDISC近期的Roadmap如下：

持續的魯棒性及效能改進

x86後端補齊計算密集型運算元的支援，端到端完整開源x86後端的支援

GPGPU上基於Stitching的大顆粒度自動程式碼生成

AMD rocm GPU後端的支援

PyTorch訓練場景的支援

此外，在中長期，我們在下面幾個探索性的方向上會持續投入精力，也歡迎各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社群建設感興趣的同行一起參與共建。

更多新硬體體系結構的支援和適配，以及新硬體體系結構下軟硬體協同方法學的沉澱

計算密集型運算元自動程式碼生成和動態shape語義下全域性layout最佳化的探索

稀疏子圖的最佳化探索

動態shape語義下執行時排程策略、記憶體/視訊記憶體最佳化等方面的探索

模型壓縮與編譯最佳化聯合的技術探索

圖神經網路等更多AI作業型別的支援和最佳化等

開源地址：

https://github.com/alibaba/BladeDISC

參考文獻

"DISC: A Dynamic Shape Compiler for Machine Learning Workloads", Kai Zhu, Wenyi Zhao, Zhen Zheng, Tianyou Guo, Pengzhan Zhao, Feiwen Zhu, Junjie Bai, Jun Yang, Xiaoyong Liu, Lansong Diao, Wei Lin
Presentations on MLIR Developers' Weekly Conference: 1, 2

"AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures", Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Jidong Zhai, Shuaiwen Leon Song, and Wei Lin. The 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2022. [to appear]
"FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads", Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang, and Wei Lin. arXiv preprint