編譯器與IR的思考: LLVM IR,SPIR-V到MLIR
推薦語
2017年,我曾經(jīng)寫過一篇文章“Deep Learning的IR“之爭”“,彼時機器學習編譯器(ML Compiler)的實踐剛開始不久。經(jīng)過幾年的探索和演進,ML Compiler已經(jīng)成為ML System中最重要的一環(huán)。ML Compiler既有其自身的特點,又和傳統(tǒng)編譯器有著密切的聯(lián)系;既有新的挑戰(zhàn),又有很多前人的經(jīng)驗可以借鑒,是個有趣且有高價值的話題。本文作者除了直接參與ML Compiler研發(fā),也對很多相關問題有深入的思考,之前的很多交流也給我很多啟發(fā),非常高興能邀請他來和大家討論這個話題。?本文為系列文章的第一篇,篇幅較長但內(nèi)容非常精彩,推薦大家完整閱讀。更習慣閱讀英文版本的朋友,可以點擊文末原文鏈接訪問。唐杉
作者簡介?
張磊,畢業(yè)于北京大學,現(xiàn)就職于谷歌。一直從事編譯器與軟件工具鏈開發(fā),涉及領域包括圖形渲染以及人工智能。SPIR-V和MLIR開發(fā)組早期核心成員之一,給SPIR-V和MLIR基礎工具鏈做過大量開源貢獻,現(xiàn)主要主導推動基于Vulkan標準和SPIR-V編譯的途徑來利用任意GPU來運行機器學習任務。作者序
編譯器 (Compiler) 通常是各種提高開發(fā)效率的軟件工具鏈中不可或缺的部分。編譯器一般被認為是黑箱,吃進高抽象層次的源程序,產(chǎn)生語義不變,可在目標硬件上運行的低層次機器碼。當然,編譯器也有其內(nèi)部結構,中間表示 (IR: Intermediate Representation) 串聯(lián)起編譯器內(nèi)各層級和模塊。
中間表示對編譯器至關重要,也如編譯器一樣百花齊放。我在日常工作中有幸能夠涉及三種主流編譯器中間表示或者基礎設施——LLVM IR, SPIR-V, 以及 MLIR, 尤其對于后兩種,我都參與了早期的開發(fā)。我打算用一系列文章記錄自己對于編譯器以及中間表示的理解,希望對感興趣的人有所幫助。
這是開坑第一篇,主要討論我個人對于編譯器以及中間表示的演進歷史以及發(fā)展趨勢的理解。著眼點在于為什么現(xiàn)有中間表示會設計成它們當前的形態(tài),而非具體的機制以及如何使用。后者已經(jīng)有詳盡的規(guī)范和教程可以參考。話題所致,此篇會比較抽象,或者用最近比較拉風的詞,形而上 (meta)
,在后續(xù)的文章中我會更加具體。另外,此篇也會盡量使用比較和類比,因為一般而言,借助于已熟悉概念來理解新概念會比較高效。討論偶爾也會涉及非編譯器的領域,但保證它們是相關的,一法通,萬法通。不過,盡管我不會離題太久,依然可能比較適合爆米花食用。閑言少述——
在討論各種具體中間表示之前,先讓我們總體看一下編譯器和中間表示。
抽象與語義(Abstractions and Semantics)
自人類文明產(chǎn)生至今,雖然科技演進的速度飛快,人腦卻沒有多少變化。我們理解不斷爆炸式提升的復雜度的方式是抽象 (abstraction)。抽象幫助我們忽略細枝末節(jié),著眼于主要矛盾。抽象減少了人腦需要處理的變量數(shù)量,減輕了人腦的處理負擔,這對于求解復雜問題至關重要。自抽象而產(chǎn)生的問題描述被稱為模型 (model),例如,在編程語言中,我們有機器模型 (machine model) 用以描述底層計算架構;在數(shù)據(jù)系統(tǒng)中,我們有數(shù)據(jù)模型 (data model) 用以描述數(shù)據(jù)關系結構;在分布式系統(tǒng)中,我們有系統(tǒng)模型 (system model) 用以描述時序和容錯假設。模型通常是原體的理想化描述。有時模型會離現(xiàn)實情況非常遠,例如,估計沒人能夠設計準確的股市模型。但無論如何,模型始終是保持復雜度可控不可或缺的工具——模型能夠給我們清晰明確的語義 (semantics), 即因定義而一定成立的原理。有了這些,我們才能對模型所描述的原體進行邏輯甚至數(shù)學層次的嚴謹推理 (reasoning)。人腦本質(zhì)性喜歡邏輯與解釋(神話、科學、甚至迷信都是人類試圖解釋世界的方式)。此處似乎離題了,實際上我只是在用廣義的方式引入一些常見的編譯器開發(fā)術語——抽象、模型、語義、推理。當然,計算機可行的其它領域也會關注這些方面,比如,我們希望面向?qū)ο蟮摹?span style="color:rgb(84,84,84);font-size:14px;">類”有良好的抽象,但這些領域通常多是“設計藝術”。相對而言,編譯器更多關注理論科學的方面;編譯器需要明確的語義用來證明 (prove) 代碼轉(zhuǎn)換的正確性。這就引入了——正確性與優(yōu)化(Correctness and Optimization)
編譯器的首要任務是保證轉(zhuǎn)換的正確性 (correctness);優(yōu)化 (optimization) 是次要的考慮。產(chǎn)生不符合源程序意圖的非?!皟?yōu)化”的代碼沒有任何意義。在整個編譯流程中正確性都需要得到保障。離開清晰的語義,正確性就無法得到定義和保證。因此編譯器對語義未知的操作 (operation) 無法進行任何轉(zhuǎn)換。當然,這并不包括每一個微小轉(zhuǎn)換,編譯器內(nèi)部也有不同層次的“邊界”。每個轉(zhuǎn)換遍歷 (pass) 保持原子性;在其內(nèi)部,可能會臨時違反源程序語義,但在每個轉(zhuǎn)換遍歷之后,中間表示應該是正確的。編譯器極度依賴每個遍歷之后的中間表示驗證 (validation) 來保證正確性。在正確性得到保障之后,我們才可以考慮優(yōu)化。生成高性能代碼似乎是無可爭議的目標,但實際上也充滿了微妙的細節(jié)。不同的源程序有著不同的模式,不同的硬件喜歡不同的指令流。編譯器處于源程序和目標硬件之間,實際上能夠動用來提升性能的手段是非常有限的。編譯器需要在大量的限制因素下做抉擇,并且保證這些抉擇對大多數(shù)場景適用,特別是針對 C++ 這種通用編程語言。這其實是非常高的要求,導致最后只有少數(shù)一些轉(zhuǎn)換(像是不可達代碼消除DCE、常量折疊constant folding、指令標準化canonicalization等等)可以全局默認開啟,很多其他轉(zhuǎn)換只能在之后的第二階段,針對目標硬件優(yōu)化時使用。在編譯器一步步轉(zhuǎn)換程序的過程中,越來越多的高層次的簡明的信息被打散,轉(zhuǎn)換成低層次的細碎的指令,這個過程被稱為代碼表示的遞降?(lowering)(不知道公認的中文翻譯是啥,只能自己來了。
);與之相反的過程被稱為代碼表示遞升?(raising)。后者通常遠比前者困難,因為后者需要在蕪雜的細節(jié)中找出宏觀脈絡。代碼表示遞降是編譯器的通常轉(zhuǎn)換方式。不難理解,越晚執(zhí)行的轉(zhuǎn)換越有結構性劣勢,原因是缺乏高層次信息。這限制了目標硬件相關優(yōu)化所能解決的問題也決定了實現(xiàn)的復雜程度。其實這里的本質(zhì)性問題是強耦合 (coupling)。這種強耦合綁定了不同應用領域(尤其是使用通用編程語言的情況)和編譯器中不同垂直轉(zhuǎn)換路徑(尤其是編譯器使用統(tǒng)一IR的情況)。解耦 (decoupling) 是我們在各種領域中實現(xiàn)更加復雜系統(tǒng)和支持更加高級場景的一般方式??梢灶A見,領域?qū)S谜Z言以及編譯器內(nèi)部解耦能夠更好地釋放編譯器的優(yōu)化能力。在接下來的 MLIR 章節(jié)我會進一步展開這一點。總而言之,成熟的編譯器可能會幫助其支持的各種場景挖掘出大部分的性能,比如說是 80%,但期待編譯器對所有應用達到最優(yōu)性能是不切實際的。編譯器的真正優(yōu)勢是幫助我們節(jié)省達到前 80% 性能所需的工程投入,讓我們可以集中資源在剩下的 20% 或者其他核心問題上。效率工具
盡管已經(jīng)近乎陳詞濫調(diào),還是再一次強調(diào),編譯器只是提升開發(fā)效率的工具。寫匯編碼或者機器碼也許會顯得特別酷,卻很難是高效的開發(fā)方式(不過,對于某些適合的場景,通過手寫匯編來獲取極限性能卻是非常合理和常見的方式)。 能夠使用高層次抽象的語言可以讓開發(fā)者忽略底層芯片上具體的寄存器和指令,從而避免陷入耗時易錯的繁雜細節(jié)之中,是對開發(fā)效率的極大提升。這就涉及我們管理不斷提升的復雜度的方式——我們有了針對不同層次的抽象,只需要創(chuàng)建工具來對這些抽象進行自動轉(zhuǎn)換。當然,并非所有的轉(zhuǎn)換都可以自動化。對于可行的,我們都可以把這種工具看作廣義上的編譯器。比如,在數(shù)據(jù)處理系統(tǒng)中,Apache Beam 提供了統(tǒng)一的語言來描述任務。Beam 會把這些任務轉(zhuǎn)換到具體的執(zhí)行引擎 (Spark, Flink, 或者其他) 的描述。這些執(zhí)行引擎再進一步把這些任務編譯成運行在具體機器上的操作并負責調(diào)度。這里的整個流程也可以看做是一種編譯。類編譯器的工具通過隱藏底層細節(jié)、提供高層次抽象而極大地提升了開發(fā)者的效率。編譯器通過對中間表示進行一系列變換 (transformation) 來鏈接不同層次的抽象。IR的形態(tài)和兼容性
中間表示只是程序的一種表示,其設計注重支持變換操作,使其正確和高效。當然,后面我們也會發(fā)現(xiàn)中間表示也會實現(xiàn)一些其他目標。在早期,編譯器通常有單一中間表示,但隨著編譯器的演進,情況變得更加復雜?,F(xiàn)代編譯器通常會有多層級的內(nèi)部表示。比如,用 Clang 來編譯 C++ 程序通常要經(jīng)過 Clang AST, LLVM IR, MachineInstr, MC 等等多個層級[1]。針對 GPU,我們也會發(fā)現(xiàn)完整編譯器被拆分成離線 (offline) 和在線 (online) 兩部分。兩部分之間的程序表示,像是 SPIR-V,也會被稱為中間表示,但其實已經(jīng)不再局限于編譯器內(nèi)部了。中間表示可以有三種形態(tài):用以高效分析和變換的內(nèi)存表示 (in-memory form),用以存儲和交換的字節(jié)碼 (bytecode form),以及用以閱讀和糾錯的文本表示 (textural form)?;谑褂脠鼍?,我們通常會見到不同的設計折中來側(cè)重不同的形態(tài)和提供不同程度的兼容性。比如,LLVM IR 更側(cè)重于編譯器內(nèi)部使用,所以它的核心是高效的內(nèi)存表示并提供相對較弱的兼容性以便于迭代改進。而 SPIR-V 則設計為硬件驅(qū)動的輸入程序表示,所以它更注重于字節(jié)碼的高效處理并提供強兼容性。此處無所謂對與錯,只是滿足不同的需求而已。但這些設計折衷依然對整個上層的生態(tài)系統(tǒng)產(chǎn)生深遠影響。IR的設計理念
現(xiàn)實中并不存在中間表示的普適設計規(guī)則。大部分情況下,中間表示的設計都是在各種限制條件下權衡各種利弊然后做出折中選擇。這些選擇會考慮具體問題的普遍性或者特殊性、給整個編譯器棧帶來的組合復雜性、對各種變換的影響等等??傮w上,我們希望:- 中間表示中的操作 (operation) 能夠具有清晰明確的語義。這是一切的根基。
- 在這之上我們希望操作能夠相互正交 (orthogonal),這有助于定義標準 (canonical) 形態(tài)以減少變換需要考慮的情況。
- 另外我們也會希望在同一中間表示的不同部分中避免出現(xiàn)重復信息。重復信息在各種變換后有很大概率會導致中間表示的不一致,從而導致錯誤編譯。
- 盡可能地保持高層次信息也是非常有好處的,因為在代碼表示遞降后想重新找回丟失的信息很難。
- 以及其他種種。
- 如果硬件實現(xiàn)了某特殊的功能模塊,我們希望能夠提供相應的軟件和中間表示抽象, 即便這會“破壞”中間表示的正交性。舉個栗子,在 GPU 中間表示中,除了基礎的乘和加指令,都有 FMA 指令。
- 在 SPIR-V 中,一個編譯單元 (module) 會提前聲明所需要的硬件能力 (capability)。這些信息完全是可以通過對中間表示的主體進行分析來得到的,所以是重復的信息。但是它們的存在減少了硬件驅(qū)動所需做的工作,從而加速了實際運行。
- 如果輸入語言已經(jīng)足夠低層次,比如 C,那么我們沒有任何辦法,只能想辦法提升抽象的層次來產(chǎn)生更好的代碼,像是對標量代碼的自動向量化 (vectorization)。
LLVM IR
LLVM 最初發(fā)布于 2003 年。在經(jīng)過了接近二十年的開發(fā)之后,LLVM 技術棧已經(jīng)非常成熟,并有一個極好的生態(tài)系統(tǒng)。LLVM 支持許多前端語言和后端硬件,許多軟硬件廠商有衍生版 LLVM 來針對自己產(chǎn)品的提供各種附加功能。我想 LLVM 對于整個業(yè)界的重要性無需我再贅言。
解綁和模塊化編譯器(Decoupling and modularizing compilers)
LLVM 帶來的最重要的東西是對編譯器解綁和模塊化的實踐,由此誕生了大量優(yōu)秀的編譯器庫和工具。在前 LLVM 時代,編譯器通常是特殊設計并高度耦合。這些編譯器雖然也分為三段——前端 (frontend) 用來對源語言進行解析,中端用來進行優(yōu)化,后端 (backend) 用來產(chǎn)生機器碼,但它們一般只針對某一特定語言(含衍生語言)或者目標硬件,編譯器內(nèi)部各個模塊也沒有明確界限。不同的編譯器?;静还灿么a,我們無法組合不同編譯器棧中的現(xiàn)存的前端或者后端,從而無法真正發(fā)揮三段式編譯器的優(yōu)勢。LLVM 依靠解綁帶來巨大變革。LLVM IR [2]顯而易見地處于 LLVM 生態(tài)的核心地位。LLVM IR 使用控制流 (control flow)、基礎塊 (basic block)、以及靜態(tài)單賦值 (SSA) 形式來表示程序。這種表示是完備的,LLVM 從而可以獨立于其他表示形態(tài),實現(xiàn)作為前后端之間的單一橋梁。這就完全地解綁了編譯器的前后端。[3]這之后我們所需的只是遵循模塊化的最佳實踐。LLVM 的代碼是組織成一系列庫 (library) 的。庫的組織形式當然有其問題,但確實是經(jīng)過實踐檢驗的系統(tǒng)級模塊化方式。庫定義了不同模塊之間的分界。通過合適的編程接口 (API),我們可以選擇并且組合不同的編譯器功能來完成不同的任務, 像是通過調(diào)用 Clang 功能庫實現(xiàn)靜態(tài)分析以及代碼格式化。這些都是非常有用的。文本IR形式(Textural IR forms)
除卻解綁和模塊化,LLVM IR 還帶來了許多其他易用性和開發(fā)效率上的提升。在內(nèi)存表示之外,對文本表示的原生支持將傳統(tǒng)的 UNIX 哲學帶入了編譯器。UNIX 哲學講求每個工具負責一個簡單任務,然后利用文件 (file) 以及管道 (pipe) 串聯(lián)不同工具來實現(xiàn)復雜功能。在 UNIX 系統(tǒng)中,文件是對資源的統(tǒng)一抽象。尤其是文本文件,基本是大多數(shù)工具的信息交換媒介。文本文件足夠靈活強大,能夠支持不同的需求,同時又直觀易用。在一個處理流程中(例如cat一體兩面
LLVM 是編譯器開發(fā)的大幅躍進。其良好的設計和活躍的社區(qū)帶來了許多提升效率的工具。但凡事都有一體兩面?;诂F(xiàn)有的 LLVM 生態(tài),我們可以越來越明顯地看到有些設計折中帶來的影響。中心化和各種衍生
LLVM IR 在整個 LLVM 生態(tài)中處于中心地位,這是編譯器前后端解綁的基礎,但這也意味著完整的編譯路徑必須通過 LLVM IR。因為 LLVM IR 的核心地位,對其進行修改需要滿足極高的條件。各種工具都基于 LLVM IR,各種公司或者組織的內(nèi)部流程都會通過它實現(xiàn)。即便你無需頻繁地打通整個路徑,對 LLVM IR 的局部小修改也會帶來意想不到的間接效應。自然而然,這就意味著改動緩慢,需要長時間高強度的討論,以及各利益相關者的同意。這對于保持 LLVM IR 本身的質(zhì)量是必須的,但如果我只是有一個非常特殊的需求, 就很難勸服整個生態(tài)做出相應的修改。一種常見的方式是分裂 LLVM,把修改保持在本地,創(chuàng)建衍生版。這依然會有很高的工程代價。LLVM IR 的中心地位對把所有改動貢獻回上游更加友好。LLVM 的代碼庫每天有將近一百次提交,這些提交添加各種新功能或者修復各種問題。如果不及時追蹤上游的提交,衍生版會偏離的越來越嚴重,最后可能無法再合并。另一方面,持續(xù)不斷地追蹤則意味著專門的人力和資源投入,很多小機構無法一直負擔。總而言之,后果就是在現(xiàn)實中我們有各式各樣的 LLVM 衍生代碼,這些代碼追蹤著不同的上游版本, 有著不同的新鮮度。如果把全球作為一個整體來看,維持和更新這些衍生代碼消耗了大量的人力和資源。當然這并不能說是 LLVM 獨有的問題。以開源模式開發(fā)的大規(guī)模復雜系統(tǒng)在被各種組織商業(yè)化的時候都有類似問題。但類似項目通常會在設計時就考慮本地定制化的需求。相較而言,LLVM IR 的絕對中心地位使得其很難被本地定制化。某種意義上講,這就是一種強耦合;MLIR 就在解耦上再進一步。演進與兼容性
LLVM IR 的另一設計是能夠協(xié)同演進中間表示和各種編譯器分析以及變換。這對于工具鏈本身質(zhì)量的提升是至關重要的,但它同時帶來了相對較弱的兼容性保證[4]。當然社區(qū)不會無緣無故引入不兼容的變動,只是這一可能性確實存在。編譯器存在于近乎操作系統(tǒng)的層級,所以很容易理解人們?yōu)槭裁词褂?LLVM IR 來作為程序的表示形式傳送給硬件驅(qū)動,特別是考慮到 LLVM IR 極好的生態(tài)系統(tǒng)以及原生的字節(jié)碼支持。但是 LLVM IR 真正的使用場景是作為不同軟件模塊之間的程序表示;涉及硬件和驅(qū)動則完全是另一個故事。硬件設備會存在于像是手機之類的終端產(chǎn)品之中,這些終端產(chǎn)品被產(chǎn)品制造商和最終消費者所掌控。驅(qū)動的升級是無法得到保障的,因此驅(qū)動依賴的 LLVM 庫也可能永遠無法得到升級。
實際上已經(jīng)有許多如此使用 LLVM IR 的嘗試,有成功,也有失敗。一個比較突出的例子是 Standard Portable Intermediate Representation (SPIR)。SPIR 是為表示 OpenCL 設備程序 (kernel) 而設計的,它鎖定了某一版本的 LLVM IR,使用 LLVM 內(nèi)聯(lián)函數(shù) (intrinsic) 和元數(shù)據(jù) (metadata) 來定義 OpenCL 的計算原語以及定義。但 Khronos Group 逐漸意識到 LLVM IR 實在不適合這種任務,遂轉(zhuǎn)向了設計與開發(fā) SPIR-V。
SPIR-V
SPIR-V 最初發(fā)布于 2015 年。SPIR-V 是多個 Khronos API 共用的中間語言,包括 Vulkan, OpenGL, 以及 OpenCL。定義全新的中間表示并且開發(fā)整套工具鏈需要大量的工作,但對 SPIR-V 的投入依然被認為是值得的。
標準規(guī)范的擴展性和兼容性
Khronos Group 的標語是“連接軟件與硬件”,簡明扼要地總結了它的任務。這種連接是通過標準規(guī)范 (standard) 和編程接口。Khronos Group 定義標準規(guī)范以及編程接口;硬件廠商提供它們的硬件實現(xiàn),軟件廠商則可以讓軟件在所有支持的平臺與設備上運行。Khronos Group 定義維護了很多標準規(guī)范,比較著名的有 Vulkan, OpenGL, 以及 OpenCL。標準規(guī)范的主要目的是抽象不同的硬件實現(xiàn),并提供對上層軟件的統(tǒng)一接口。但同時,標準規(guī)范也需要能夠支持硬件特有的功能。這是對現(xiàn)實中存在各式各樣硬件的一種承認,也讓軟件能夠深度挖掘某些具體硬件的性能。這兩個看似相互沖突的目標通過分等級的特性 (feature) 體系 (hierarchy) 得以支持。Khronos Group 內(nèi)部有清晰的流程來管理各種特性,包括提議新的特性,提升某廠商專有的特性為通用特性等等。SPIR-V 也是一種標準,所以具有相同的設置。除了核心特性之外,SPIR-V 支持通過多種機制來擴展其功能,包括添加新的枚舉值, 引入新的擴展 (extension),或者通過某個命名空間引入一整套指令 (extended instruction set)。其擴展也分為不同等級——廠商自有擴展 (vendor specific)、多廠商聯(lián)合支持的擴展 (EXT)、 以及 Khronos 級別的擴展 (KHR)。任意廠商都可以提議新的擴展,但擴展越接近核心級別,就需要越多的廠商復議,并且要經(jīng)過愈發(fā)嚴格的審核和批準流程??傮w上 SPIR-V 和其工作組 (working group) 分別從技術和組織上提供了框架來支持標準的演進和擴展。[5]LLVM IR 也提供一些方式來擴展其中間表示,尤其是內(nèi)聯(lián)函數(shù) (intrinsic) 和元數(shù)據(jù) (metadata),但難以想象其支持所有廠商的各種自有內(nèi)聯(lián)函數(shù),更不用說引入各種廠商自有的類型以及指令模式到核心 LLVM IR 指令中。除此之外,連接軟硬件的標準規(guī)范也非常注重穩(wěn)定性和兼容性,因為硬件驅(qū)動的更新頻度遠低于軟件工具鏈。驅(qū)動在設備的生命周期內(nèi)即便是永遠得不到更新也并不稀奇。比如,許多在低端 Android 手機上的 Vulkan 驅(qū)動停留在發(fā)布于 7 年前的1.0 版本。如果我們使用 LLVM IR 作為設備程序表示 (kernel),則有很大的可能程序會由近期 LLVM 版本的工具鏈產(chǎn)生,但是驅(qū)動卻依然停留在很早期的版本。這種版本錯配會造成各種各樣的麻煩問題。相對而言,SPIR-V 則通過版本和擴展機制以及穩(wěn)定的字節(jié)碼提供了必須的穩(wěn)定性和兼容性。穩(wěn)定的字節(jié)碼(Stable binary form)
完整的 GPU 編譯器被分為兩部分——首先通過離線工具鏈從高層次源代碼生成 SPIR-V,然后通過驅(qū)動內(nèi)部編譯器將 SPIR-V 在線編譯成機器碼。雖然像 LLVM IR 一樣在整個編譯流程中處于“中間”位置,SPIR-V 更側(cè)重于驅(qū)動內(nèi)部二次編譯的高效,因為這一步在運行時進行。所以 SPIR-V 的核心是其字節(jié)碼。其編碼有很多簡化驅(qū)動二次編譯的考量,像是用各種提前的顯示聲明來避免運行時復雜的分析。SPIR-V 并沒有在規(guī)范中指定內(nèi)存表示或者文本表示,這些都是實現(xiàn) SPIR-V 標準規(guī)范的工具鏈自行定義的。比如 SPIRV-Tools 有其自己的內(nèi)存表示和文本表示, 同樣 MLIR 中的 SPIR-V dialect 也是。GPU 領域?qū)S?/span>
至此我解釋了 SPIR-V 的 standard 和 portable,卻還并未涉及其中 IR 的部分。其實 SPIR-V 的 IR 部分和 LLVM IR 相差并不太大。SPIR-V 借鑒了很多 LLVM IR 的設計——它同樣是由控制流、基本塊、以及靜態(tài)單賦值來表示程序。指令的粒度和 LLVM IR 也相差不大。SPIR-V 中獨特的部分在于對很多 GPU 概念的原生支持。這種支持通過很多 SPIR-V 獨有的機制來實現(xiàn),比如 decorations, builtins, 以及特殊的指令(像是導數(shù)計算、圖像取樣)。另外為了支持圖形圖像和高性能計算的兩種使用場景, SPIR-V 中有許多執(zhí)行模型和模式。當然,對圖形圖像也有 structured control flow 的特殊需求。[6]一個 GPU 為主的標準規(guī)范需要原生支持各種 GPU 概念,能夠提供不同等級的擴展需求, 以及提供穩(wěn)定和兼容的字節(jié)碼。這些需求并不符合 LLVM IR 的設計理念,所以 Khronos Group 推出了 SPIR-V。但是設計一套中間表示只是個開始,圍繞其開發(fā)和維護整套工具鏈需要持續(xù)不斷的工程投入。SPIR-V 與 LLVM IR 完全無關,SPIR-V 的編譯器棧無法利用現(xiàn)有的 LLVM 庫。所以 SPIR-V 的整個棧是從頭開始獨立開發(fā)的,從匯編、反匯編,一步步到各種語言的編譯器和優(yōu)化。如果我們當時能夠有一套幫助開發(fā)編譯器的基礎設施——MLIR
MLIR 在 2019 年底合并到了 LLVM 代碼庫,所以它才開源開發(fā)了兩年左右。我個人感覺一個生態(tài)至少需要五年才能相對完善。從這個角度,MLIR 尚處于非常早期,許多開發(fā)尚待進行。但從目前的發(fā)展來看,MLIR給編譯器開發(fā)帶來了許多新穎的想法和深遠的改變,比如通過基礎設施化來進一步解耦編譯器的中間表示。
基礎設施化
基礎設施化 (infrastructurization) 其實是技術演進的自然終點 (endpoint)。成為基礎設施意味著一項技術已經(jīng)足夠成熟并且得到廣泛部署。變成基礎設施后,新的技術變革在其之上產(chǎn)生。電力、網(wǎng)絡、公共云等等,都符合這個潮流。上述都是超大規(guī)模的基礎設施化。基礎設施化同樣適用于影響規(guī)模小的技術,因為它能夠通過共用來降低開發(fā)成本, 并讓每個人只關注其核心商業(yè)邏輯的實現(xiàn)。
很多人通過機器學習編譯器了解到 MLIR。這確實是 MLIR 最初關注的領域,但 MLIR 其實有著更加廣泛的應用范圍。MLIR 是用來開發(fā)編譯器基礎設施。它提供一系列可復用的易擴展的基礎組件,用來搭建領域?qū)S镁幾g器。在 LLVM IR 和 SPIR-V 中,我們有唯一的中間表示,其中含有完備的指令集來編譯所有的 CPU 和 GPU 程序。MLIR 中則沒有完全處于中心地位的中間表示。MLIR 提供基礎設施來幫助定義 operation 以及將邏輯相關的 operation 組合成 dialect。另外,MLIR 也提供一些普適的 pattern 或者 pass,這些 pattern 或者 pass 并不與具體的 operation 綁定,能夠自適應。[7]
無論是對 operation 還是 pattern/pass 的支持都要求 MLIR 以更加細的粒度看待編譯器。在 MLIR 中,operation 不再是最基礎的部件,粒度進一步細化到類型, 值, attribute, region, 以及 interface (例如 attribute/type/operation interface).[8]
Operation 可以有任意數(shù)量的輸入、輸出、attribute,并包含任意數(shù)量的 region。其中 region 能夠表示 operation 之間的嵌套關系,從而簡化編譯器的分析和轉(zhuǎn)換。Operation 可以實現(xiàn) operation interface,pattern 和 pass 綁定的是 operation interface,由此而實現(xiàn)與具體 operation 的解綁并做到自適應。
MLIR 里面的概念都設計的比較抽象,目的是能比較好地映射到不同的領域和場景。
Dialects, dialects, dialects
當然,這套基礎設施存在的目的是幫助搭建最終編譯器。我們在寫 C++ 程序的的時候會調(diào)用 STL 或者更加高層次的庫,很少會從頭開始實現(xiàn)所有的細節(jié)。另外,基礎設施也需要與其支持的領域協(xié)同發(fā)展,因為使用場景中會提供很多需求。因此,MLIR 代碼庫中自帶很多用來給各種層級概念建模的 dialect。[9]
MLIR 的 dialect 生態(tài)目前還在擴張演進階段,但 dialect 之間的組織結構以及有些 dialect 已經(jīng)相對穩(wěn)定了。比如我們有 LLVM 和 SPIR-V dialect 作為與其他系統(tǒng)轉(zhuǎn)換的邊界 dialect。(其實 MLIR 可以同時表示 LLVM IR 和 SPIR-V 這一點也表明了 MLIR 的基礎設施角色。) 抽象層次居中的有 Linalg, Tensor, Vector, SCF dialect,它們協(xié)同合作用來生成代碼。另外,MLIR 中還有 Affine, Math, Arithmetic dialect 用來描述底層計算。在 AI 框架層面,有 TensorFlow, TFLite, MHLO, Torch, TOSA 進行對接和導入模型。除此之外,還有很多其他用途的 dialect,像是 PDL 用來定義編譯器轉(zhuǎn)換等等。
Alex 之前在 MLIR 論壇上分享的各 dialect 之間的關系[10]非常值得一讀,之后我也會寫下我的理解。這些各式各樣的 dialect 和以后包裝它們而產(chǎn)生的局部或者完整的轉(zhuǎn)換流程將極大簡化領域相關編譯器的開發(fā)。
進一步解耦編譯器和中間表示
其實基礎設施化以及由此產(chǎn)生的大量 dialect 都是進一步解耦和模塊化編譯器以及中間表示的一種自然結果。唯一的中間表示被許多以 dialect 形態(tài)存在的部分的中間表示取代。沒有某個部分中間表示再處于中心地位,都是按需組合。如果現(xiàn)有的 dialect 不能滿足需求,定義新的 dialect 非常簡單。再過幾年待 MLIR 更加完善后,可以預見,開發(fā)領域?qū)S镁幾g器只需要新加自己的邊界 dialect 來定義項目相關的 operation,之后就是選取和組合現(xiàn)有 dialect 來形成整體的流程。如果這個構想實現(xiàn),那肯定會極大縮減開發(fā)所需工程投入。
另外,進一步解耦中間表示也讓我們可以靈活地根據(jù)領域進行設計和折中。我們只需選取所需的部分中間表示來組合成完整編譯器,不再需要全盤接收像 LLVM IR 一樣的一套完整中間表示。因為 interface 的存在,擴展模塊的更能也變得更加簡單——我們既可以定義新的 operation 來實現(xiàn)已有的 interface,也可以定義新的 interface 然后支持現(xiàn)有 operation。
換言之,LLVM IR 天然中心化并且偏好統(tǒng)一的編譯流程,MLIR 的基礎設施和 dialect 生態(tài)則天然是去中心化并且偏好離散的編譯流程。
技術的一般發(fā)展趨勢是從單一的強耦合整體到適用不同場景的多種多樣的選擇。對于技術棧的上層而言,這尤其明顯,因為越往上越接近用戶和商業(yè)需求,而用戶和商業(yè)需求本身就各式各樣,由層出不窮的前端框架可見一斑。
技術棧的底層一般相對穩(wěn)定。少數(shù)幾種硬件架構、編譯器和操作系統(tǒng)統(tǒng)治很多年。但半導體進展的變慢和計算需求的爆炸式增長也在驅(qū)動著底層技術的變革?,F(xiàn)在依然依靠通用架構和普適優(yōu)化很難再滿足各種需求,開發(fā)領域?qū)S玫恼w的解決方案是一條出路。RISC-V 在芯片指令集層次探索模塊化和定制化,MLIR 則是在編譯器以及中間表示層面做類似探索。兩者聯(lián)手會給底層技術棧帶來何種革新是一個值得拭目以待的事情。
跨系統(tǒng)邊界的漸進式代碼表示遞降
在結束本章之前,再啰嗦最后一點。其實我們可以從兩個維度看待 MLIR 帶來的解耦:
水平方向上,dialect 把完整中間表示打散成許多局部中間表示;垂直方向上,MLIR 讓我們可以對處于不同層級的概念進行建模。這對領域?qū)S镁幾g器是非常有用的,因為領域?qū)S谜Z言一般是高度抽象的聲明式語言,只描述任務,需要編譯器將其轉(zhuǎn)換成具體的命令式機器指令。一步跨越這個巨大的抽象差距是非常難的,利用多級抽象和建模來進行漸進式 lowering 是更加適合的方式。我們可以分離各個層次關注的問題,整個系統(tǒng)也更加的易開發(fā)和維護。
當然這并不是什么全新的概念,在不同的項目中我們已然看到各種類中間表示的設置,像是 Clang AST 或者各種機器學習框架中的計算圖。MLIR 的優(yōu)勢是使用同樣的基礎設施將這些不同層次的表示連接起來,讓它們之間的信息流通變得更加順暢。其實現(xiàn)代復雜系統(tǒng)的開發(fā)多是選取各種子系統(tǒng)然后將其組合。將來自前一個子系統(tǒng)的數(shù)據(jù)進行驗證、轉(zhuǎn)化然后傳遞給下一個子系統(tǒng)消耗掉很多工程資源。如果所有子系統(tǒng)使用相同的內(nèi)部基礎設施,這些資源投入就都可以節(jié)省下來,另外,使用相同工具也會使得跨組跨項目的溝通協(xié)調(diào)變得更加簡單。
結語:
參考資料
[1]?“clang c++ compilation涉及的步驟 ”,https://eli.thegreenplace.net/2012/11/24/life-of-an-instruction-in-llvm?
[2]?“LLVM語言參考”,https://llvm.org/docs/LangRef.html?
[3]?“LLVM解耦的部分”,https://www.aosabook.org/en/llvm.html?
[4]?“LLVM兼容性的部分”,https://llvm.org/docs/DeveloperPolicy.html#ir-backwards-compatibility?
[5]?“SPIR Ecosystem”,https://www.khronos.org/spir/?
[6]?“SPIR-V的機制”,https://www.khronos.org/registry/SPIR-V/specs/unified1/SPIRV.html
[7]?“MLIR語言參考”,https://mlir.llvm.org/docs/LangRef
[8]?“MLIR Interface”,https://mlir.llvm.org/docs/Interfaces/
[9]?“MLIR Dialects”,https://mlir.llvm.org/docs/Dialects/
[10]?“MLIR dialect overview”,https://discourse.llvm.org/t/codegen-dialect-overview/2723
