Google全面開源 MLIR 及生態聯盟,全球 95% 的加速器硬體都在使用

  • 2019 年 10 月 4 日
  • 筆記

近日Google宣布,向非盈利性 LLVM 基金會提供今年 4 月開源的 Multi-Level Intermediate Representation(MLIR)架構,一個與 TensorFlow 緊密結合的表示格式和編譯器實用工具庫,該架構介於模型表示和低級編譯器/執行器(二者皆可生成硬體特定程式碼)之間。Google希望通過向社會提供該架構來激勵更多的創新,從而進一步加速 AI 領域發展。我們將 MLIR 的詳細內容及Google相關報道整理編譯如下。

MLIR 與 TensorFlow 的淵源


在過去,若想解決多級別堆棧問題,則需要我們構建新的軟硬體堆棧生成器,這也意味著必須為每個新路徑重新構建優化與轉換傳遞。

TensorFlow 生態系統包含許多編譯器和優化器,可在多個級別的軟硬體堆棧上運行。作為 TensorFlow 的日常用戶,在使用不同種類的硬體(GPU、TPU、移動設備)時,這種多級別堆棧可能會表現出令人費解的編譯器和運行時錯誤。

圖 1 TensorFlow 組件概述

TensorFlow 能夠以多種不同的方式運行,如:

  • 將其發送至調用手寫運算內核的 TensorFlow 執行器
  • 將圖轉化為 XLA 高級優化器(XLA HLO)表示,反之,這種表示亦可調用適合 CPU 或 GPU 的 LLVM 編輯器,或者繼續使用適合 TPU 的 XLA。(或者將二者結合)
  • 將圖轉化為 TensorRT、nGraph 或另一種適合特定硬體指令集的編譯器格式
  • 將圖轉化為 TensorFlow Lite 格式,然後在 TensorFlow Lite 運行時內部執行此圖,或者通過 Android 神經網路 API(NNAPI)或相關技術將其進一步轉化,以在 GPU 或 DSP 上運行

但事實上,多級別堆棧的複雜性遠遠超過圖 1 所示。為了更好解決 TensorFlow 用戶在使用不同種類的硬體(GPU、TPU、移動設備)時,由於多級別堆棧而導致的編譯器與運行時錯誤,我們開源了一個全新的中介碼與編譯器框架 MLIR。

什麼是 MLIR


MLIR(或稱為多級別中介碼)是一種表示格式和編譯器實用工具庫,介於模型表示和低級編譯器/執行器(二者皆可生成硬體特定程式碼)之間,在生產品質組件的支援下,能夠對優化編譯器設計與實現進行全新探索。

圖 2 Google MLIR(相關 ppt 見文末)

MLIR 深受 LLVM 的影響,並不折不扣地重用其許多優秀理念,比如擁有靈活的類型系統,可在同一編譯單元中表示、分析和轉換結合多層抽象的圖等。這些抽象包括 TensorFlow 運算、嵌套的多面循環區域乃至 LLVM 指令和固定的硬體操作及類型。

為區分不同的硬體與軟體受眾,MLIR 提供「方言」,其中包括:

  • TensorFlow IR,代表 TensorFlow 圖中可能存在的一切
  • XLA HLO IR,旨在利用 XLA 的編譯功能(輸出到 TPU 等)
  • 實驗性仿射方言,側重於多面表示與優化
  • LLVM IR,與 LLVM 自我表示之間存在 1:1 映射,可使 MLIR 通過 LLVM 發出 GPU 與 CPU 程式碼
  • TensorFlow Lite,將會轉換以在移動平台上運行程式碼

每種方言均由一組存在不變性的已定義操作組成,如:「這是一個二進位運算符,輸入與輸出擁有相同類型。」

MLIR 沒有眾所周知的固定或內置的操作列表(無「內聯函數」),方言可完全定義自定義類型,即 MLIR 如何對 LLVM IR 類型系統(擁有一流匯總)、域抽象(對量化類型等經機器學習 (ML) 優化的加速器有著重要意義),乃至未來的 Swift 或 Clang 類型系統(圍繞 Swift 或 Clang 聲明節點而構建)進行建模。

另外值得一提的是,雖然 MLIR 充當 ML 的編譯器,但它同樣支援在編譯器內部使用機器學習技術。MLIR 的擴展性有助於探索程式碼降階策略,並在抽象之間執行逐步降階。

MLIR 開放的意義


機器學習現在的使用範圍非常廣泛,它可以在從包含 GPU 和 TPU 的雲基礎設施到行動電話,甚至是最小的硬體(例如為智慧設備供電的微控制器)上運行。正是因為將硬體和開源軟體框架(如:TensorFlow)的優勢相結合,今天我們才能看到所有令人難以置信的 AI 應用成為可能。無論是預測極端天氣(https://www.youtube.com/watch?v=p45kQklIsd4);幫助有語言障礙的人更好地溝通;還是協助農民檢測農作物疾病。

圖 3 AI 協助農民進行檢測(https://www.blog.google/technology/ai/ai-takes-root-helping-farmers-identity-diseased-plants/)

但隨著所有這些進展如此迅速,企業工廠方正在努力跟上不同的機器學習軟體框架與各種不斷增長的硬體組合。機器學習生態系統依賴於許多不同的技術,而這些技術通常具有不同的複雜程度,因而無法很好地協同工作。

管理這種複雜性的負擔最終落在了研究人員、企業和開發人員身上。通過減緩新的機器學習驅動產品從研究到實現的速度,這種複雜性將會影響我們解決具有挑戰性現實問題的能力。

今年早些時候,我們發布了 MLIR,這是一種開源機器學習編譯器基礎架構,可以解決因軟體和硬體碎片不斷增加而導致的複雜性,並且可以更輕鬆地構建 AI 應用程式。它提供了新的基礎設施和設計理念,使得機器學習模型能夠在任何類型的硬體上一致地表示和執行。現在,我們宣布我們會向非營利性 LLVM 基金會(http://llvm.org/foundation/)提供 MLIR ,這也將使整個行業更快地採用 MLIR。

圖 4 MLIR 生態聯盟

MLIR 旨在成為 ML 基礎架構的新標準,並得到全球硬體和軟體合作夥伴的大力支援,包括 AMD,ARM,Cerebras,Graphcore,Habana,IBM,Intel,Mediatek,NVIDIA,Qualcomm Technologies,Inc,SambaNova Systems,Samsung,Xilinx 的小米——佔全球數據中心加速器硬體的 95%以上,超過 40 億部手機和無數的物聯網設備。在Google,MLIR 正在整合併用於我們所有的伺服器和移動硬體工作。

機器學習已經走過了漫長的道路,但之後的路仍然很長。通過 MLIR,人工智慧將通過賦予研究人員更大規模地訓練和部署模型的能力,以及在不同硬體上具有更高的一致性、速度和簡單性,從而更快地推進該領域的發展。這些創新也可以迅速進入你每天使用的產品中,並在你的所有設備上順利運行。我們也希望通過 MLIR 能夠最終實現 AI 對地球上的每個人都更有幫助、更有用的願望。

關於 MLIR 的 ppt: http://llvm.org/devmtg/2019-04/slides/Keynote-ShpeismanLattner-MLIR.pdf 文章相關鏈接: https://medium.com/tensorflow/mlir-a-new-intermediate-representation-and-compiler-framework-beba999ed18d https://www.blog.google/technology/ai/mlir-accelerating-ai-open-source-infrastructure/ MLIR 開源地址: https://github.com/tensorflow/mlir