1. 引言
過去的幾年裡推動機器學習技術穩步發展的根本性改變之一是訓練和優化機器學習模型的巨大計算力。許多技術都是很年前就已經提出，唯有近幾年提升的計算力可以為現實世界的問題提供足夠優質的解決方案。這些計算能力的很大一部分是通過 GPU 獲取的，其針對向量的計算能力最初是為圖形而設計的，但機器學習模型通常需要執行復雜的矩陣運算，因此 GPU 同樣表現出了非常好的效能。
這些方法及 GPU 在現實世界，尤其是在機器學習領域的成功引發了硬體設計者的一系列創新，他們致力於為機器學習工作負載研發新的加速器。然而，儘管 GPU 很長一段時間都在 CUDA 等軟體系統發力，但這些庫通常不會擴充套件到新的非 GPU 加速器，為這些加速器開發軟體仍然是一大挑戰。
2017 年，谷歌宣佈他們將通過雲服務向大眾提供他們專有的 TPU 機器學習加速器。最初，TPU 的使用侷限於根據谷歌 TensorFlow 機器學習框架編寫的應用。幸運的是，2018 年 9 月，谷歌通過底層 XLA（Accelerated Linear Algebra）編譯器的 IR 開放了 TPU 的訪問許可權。這個 IR 是一個通用的優化編譯器，用於表達線性代數基元的任意計算，因此為使用 TPU 的非 TensorFlow 使用者以及非機器學習工作負載提供了良好的基礎。
在本文中，我們介紹了使用這個介面編譯通用 Julia 程式碼的初步工作，它們可以進一步訪問谷歌雲的 TPU。這一方法與 TensorFlow（Abadi et al., 2016）採用的方法形成對比，後者沒有編譯 python 程式碼，而是先用 Python 構建一個計算圖，然後再對這個計算圖進行編譯。它在美學上類似於 JAX（Frostig et al., 2018），JAX 的目標是通過跟蹤和 Offload 高階陣列運算來 Offload Python 本身編寫的計算。然而重要的是，我們不依賴於追蹤，而是利用 Julia 的靜態分析和編譯能力來編譯整個程式，包括傳遞到裝置端的所有控制流。
值得一提的是，我們的方法允許使用者在編寫模型時充分利用 Julia 語言的表現力。這些表現力主要體現在一些高階特徵上，如多重派發、高階函式和現有庫，如微分方程求解器（Rackauckas & Nie，2017）和通用線性代數例程等。由於只在純 Julia 程式碼上執行，所以它也與 Zygote.jl（Innes, 2018）自動微分工具相容，該工具能執行自動微分作為高階編譯過程。總的來說，我們能夠編譯使用 Flux 機器學習框架編寫的完整機器學習模型，將模型的前向、反向傳播及訓練迴路融合成一個可執行檔案，並 Offload 到 TPU 中。
論文：Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs
論文連結：https://arxiv.org/abs/1810.09868
摘要：谷歌的雲 TPU 是一種前景廣闊的新型機器學習工作負載硬體架構，近年來已經成就了谷歌很多里程碑式的機器學習突破。如今，谷歌已經在其雲平臺上為大眾提供 TPU，最近又進一步開放，允許非 TensorFlow 前端使用。我們描述了一種通過這一新 API 及谷歌 XLA 編譯器將 Julia 程式的適當部分 Offload 到 TPU 的方法和實現。我們的方法能夠將 Julia 程式編寫的 VGG19 模型及其正向傳播完全融合到單個 TPU 可執行檔案中，以便 Offload 到裝置上。我們的方法與 Julia 程式碼上現有的基於編譯器的自動微分技術很好地結合在一起，因此也能夠自動獲得 VGG19 反向傳播並採用類似的方法將其 Offload 到 TPU。使用我們的編譯器訪問 TPU，我們能夠在 0.23 秒內完成批量為 100 張影象的 VGG19 前向傳播，而 CPU 上的原始模型則需要 52.4s。我們的實現僅需不到 1000 行的 Julia 程式碼，無需根據 TPU 對核心 Julia 編譯器或任何其他 Julia 包進行特有的更改。
5. 將 Julia 語義對映到 XLA
只要 Julia 程式是按照 XLA 基元來編寫的，我們就能將其編譯到 XLA。然而，Julia 程式不是根據晦澀難懂的 HLO 操作來編寫的，而是根據由 Julia 基本庫提供的函式和抽象來編寫的。幸運的是，Julia 使用了多重派發，使得根據 HLO 操作來表達標準庫的抽象變得容易。下面展示了幾個簡單的例子：
除了這些簡單的操作以外，我們還提供了高階陣列抽象的實現，尤其是 mapreduce 和 broadcast。依據 HLO 操作實現的 broadcast 大約有 20 行程式碼，為節省空間起見，此處不予展開，但「mapreduce」的實現非常簡單：
從上圖可以看到將任意 Julia 函式作為靜態計算運算的效果。由於 Julia 對泛型抽象的依賴，它只需指定極少數定義，就能覆蓋大量 API。具體來說，從 mapreduce 的定義中，我們可以自動得到在 base 中所定義運算（如 sum 和 prod）的降維。事實上，獲取足夠的 API 覆蓋來編譯 VGG19 模型的前向傳播和反向傳播需要不到 200 行定義。
5.1 結構對映
我們做了一個額外的識別。embedded IR 中的任意元組或 immutable 結構被對映至一個 XLA 元組，即 julia 值 1 + 2im（由兩個整數結構組成的複雜數字）將被對映至 XLA 元組 (s64[], s64[])。我們在 XLA IR 的 Julia 嵌入中儲存該結構型別，但很顯然 XLA 不瞭解 julia 型別，因此在最終的轉換步驟中這些型別被轉換成適當的元組。類似地，（julia）元組建構函式（以及 immutable 結構的建構函式）變成了 XLA 的元組構件。元組引用（immutable 結構的欄位引用）變成了 XLA 的元組引用。
5.2 處理控制流
有一個額外的複雜問題我們還沒討論：Julia 提供的命令式控制流和 XLA 提供的函式式控制流之間的語義不匹配。為了解決 if/else 控制流模組，我們在 Julia 編譯器的 SSA IR 中檢視 φ 節點，然後將這些節點作為 XLA 函式式控制流的結果（如果在同一個合併點存在多個 φ 節點，則我們構造這些節點的元組）。導致計算流分化的條件變成了函式式控制流的條件，二者之間的任意計算都可作為函式呼叫。迴圈控制流類似條件控制流的構建，我們識別控制流圖的強連線區域，將其作為迴圈的主體。
7 結果
7.2 VGG19 前向傳播
我們的第一個複雜示例是完整 VGG19 前向傳播。我們使用 Metalhead 包中的 VGG19 實現 (Mike Innes & Contributors, 2018)，它利用 Flux (Innes & Contributors, 2017) 框架將熟悉的機器學習層（卷積層、全連線層）轉換成線性代數運算。但重要的是，Flux 框架中的每個層只是一般的函式，它們可以反過來呼叫一般的線性代數運算。因此，Flux 中表達的機器學習模型（包括 VGG19）只是一般的 Julia 函式，因此能夠使用本論文介紹的方法。
我們的編譯器能夠完全推斷、offload 和融合（fuse）VGG19 的全部前向傳播。在 Julia 級別的優化之後，頂層函式的最終 IR 包括 181 個指令（每個 HloOp 都是具備適當推斷的常數靜態引數和適當形態推斷的動態引數）。每個級別計算的 HLO operands 總數是 183（多出的兩個用於嵌入中隱藏的引數指令），29 個計算一共有 361 個 HLO operands，指令數詳情見圖 3。由於我們能夠 offload 全部前向傳播計算，因此 Julia 不參與任何評估步驟，從而可以同步執行其他任務（如為下一批准備資料）。此外，得到程式碼的效能僅受限於 XLA 生成的程式碼質量，不受限於前端（效能評估見 7.4）。我們在 ImageNet 驗證集上評估了 VGG19 模型，並驗證了得到結果與原版 Metalhead 的結果相匹配，從而驗證了生成的 XLA 程式碼準確性。
7.3 VGG19 反向傳播
為了獲取反向傳播，我們利用基於 Zygote.jl 編譯器的 AD 框架 (Innes, 2018)。Zygote 在 Julia 程式碼上執行，其輸出也是 Julia 函式（適合重新匯入 Zygote 以獲取更高階的導數，也適合編譯成針對 TPU 的模型）。如下是一個具體示例：
即模型當前值和特定訓練樣本（或者訓練樣本批）所對應的導數。我們使用 sum 作為損失函式的簡單替代。意外的是，第 6 章介紹的型別推斷修改也能夠提高所有 VGG19 反向傳播的型別推斷精度。至於前向傳播，優化和未優化的指令總數如圖 1 所示。反向傳播生成的 XLA 指令明顯多於前向傳播，其最大貢獻者之一便是 Zygote 的混合模式廣播融合（mixed mode broadcast fusion）——在一個對映核心（map kernel）中同時計算前向傳播和反向傳播。由於 XLA 目前不支援來自一個對映指令的多個輸出，該函式在多個對映指令上重複執行，因此後續需要清洗 XLA 的 DCE。一般，我們的編譯過程解決了 XLA 對對映指令的處理，因為在泛型程式碼中呼叫 Julia 對映和 broadcast 函式非常普遍。
7.4 在 TPU 上進行評估
圖 2：不同批大小對應的 VGG19 前向傳播時長。Flux CPU 是 Flux master/Julia master，但不使用 XLA 編譯器。PyTorch CPU 是同一 CPU 上的相同 PyTorch 模型。FluXLA CPU 是我們的研究在 CPU 上的 xrt 實現；FluXLA TPU (total) 是端到端時間，和客戶端報告的時間一致（包括 kernel launch 開銷和從谷歌雲把資料遷移回來，注意由於額外的網路遷移，該測量結果會出現極大的變動）；FluXLA TPU (compute) 是 TPU 上的總計算時間，和雲分析器報告的時間一致（與 FluXLA TPU (total) 不同，該測量很穩定）。所有 CPU 測量基於支援 AVX512 的 Intel(R) Xeon(R) Silver 4114 CPU @ 2.20GHz CPU。可獲取高達 20 個核心，且 CPU 基準不限於單個核心（即使在實踐中，也不是所有 CPU 基準都使用並行化）。TPU 基準僅限單個 TPU 核心。所有時間至少經過 4 次執行（除了 FluXLA CPU for N=100，因為它無法在 10 分鐘內完成一次執行）。
圖 3：被編譯為 XLA 後，Metalhead.jl VGG19 前向傳播和反向傳播的指令數分解，上圖展示了未優化（Julia 前端之後）和優化指令數（XLA 優化流程之後，與 CPU 後端所用流程類似，但沒有 HLO fusion）。每個指令數被進一步拆分為實體計算中的指令（E）和所有計算中的指令數（T）。