用TPU跑Julia程式，只需不到1000行代碼

新智元報導

來源：arXiv

編輯：肖琴

【新智元導讀】將Julia代碼直接部署到谷歌Cloud TPU，讓程式運行更快的官方指南來了！Julia和TPU的結合意味著快速、易於表達的ML計算！”

Julia是一門集眾家所長的編程語言。隨著Julia 1.0在8月初正式發布，Julia語言已然成為機器學習編程的新寵。

這門由 MIT CSAIL 實驗室開發的編程語言結合了 C 語言的速度、Ruby 的靈活、Python 的通用性，以及其他各種語言的優勢於一身，並且具有開源、簡單易掌握的特點。

隨著用戶越來越多，圍繞Julia的開發工具、技術、教程等也愈加豐富。昨天，Julia開發人員Keno Fischer和Elliot Saba發表了一篇新論文Automatic Full Compilation of Julia Programs and ML Modelsto Cloud TPUs，介紹如何將Julia代碼直接部署到Google Cloud TPU，讓程式運行更快。

Jeff Dean在推特上推薦了這篇論文，評價稱：“Julia和TPU的結合意味著快速、易於表達的ML計算！”

谷歌的Cloud TPU是一種用於機器學習工作負載的很強大的新硬體架構。近年來，Cloud TPU為谷歌的許多裡程碑式的機器學習成就提供了動力。

谷歌現在已經在他們的雲平台上開放提供一般用途的TPU，並且最近已經進一步開放，允許非TensorFlow前端使用。

這篇論文描述了通過這個新的API和Google XLA編譯器，將Julia程式的適當部分解除安裝（offload）到TPU的方法和實現。

這一方法能夠將表示為Julia程式的VGG19模型的前向傳遞（forward pass）完全融合到單個TPU可執行檔案中，以便解除安裝到設備。該方法也很好地與Julia代碼上現有的基於編譯器的自動微分技術相結合，因此我們也能夠自動獲得VGG19的反向傳遞並類似地將其解除安裝到TPU。

使用這一編譯器定位TPU，能夠在0.23秒內對100張影像的VGG19前向傳遞進行評估，這與CPU上原始模型所需的52.4秒相比大幅加速了。他們的實現僅需不到1000行Julia代碼，沒有對核心Julia編譯器或任何其他Julia包進行TPU特定的更改。

具體方法和實現細節請閱讀原論文。以下主要從分別從回顧TPU硬體架構、Julia編譯器的workflow、將XLA嵌入到Julia IR的細節，以及結果與討論幾個部分進行介紹。

谷歌TPU和XLA編譯器

2017年，谷歌宣布他們將通過雲服務向公眾提供其專有的張量處理單元（TPU）機器學習加速器。最初，TPU的使用僅限於使用谷歌的TensorFlow機器學習框架編寫的應用程式。幸運的是，2018年9月，Google通過較低級別的XLA（Accelerated Linear Algebra）編譯器的IR開放了對TPU的訪問權限。該IR是通用的，是用於表示線性代數原語的任意計算的優化編譯器，因此為非Tensorflow用戶以及非機器學習工作負載的TPU目標提供了良好的基礎。

XLA（加速線性代數）是谷歌的一個部分開源編譯器項目。它具有豐富的輸入IR，用於指定多線性代數計算，並為CPU，GPU和TPU提供後端代碼生成功能。XLA的輸入IR（稱為HLO高級優化IR）在基本數據類型或其元組（但沒有元組陣列）的任意維陣列上運行。HLO操作包括基本算術運算、特殊函數、廣義線性代數運算、高級陣列運算以及用於分布式計算的原語。XLA可以執行輸入程式的語義簡化，以及執行整個程式的記憶體調度，以便有效地使用和重用可用記憶體（這是大型機器學習模型的一個非常重要的考慮因素）。

每個HLO操作都有兩種操作數：

這篇論文介紹了使用這個接口將常規的Julia代碼編譯帶TPU的初步工作。這一方法不依賴跟蹤，而是利用Julia的靜態分析和編譯功能來編譯完整的程式，包括對設備的任何控制flow。特別是，我們的方法允許用戶在編寫模型時充分利用Julia語言的完整表現力，能夠編譯使用Flux機器學習框架編寫的完整機器學習模型，將前向和後向模型傳遞以及訓練loop融合到單個可執行檔案，並將其解除安裝到TPU。

Julia編譯器的工作原理

為了理解如何將Julia代碼編譯為XLA代碼，了解常規Julia編譯器的工作原理是有益的。Julia在語義上是一種非常動態的語言。但是，在標準配置中，Julia的最終後端編譯器是LLVM（Lattner＆Adve，2004），它是一個靜態編譯器後端。

Julia編譯器需要將語言的動態語義與LLVM表示的靜態語義之間聯繫起來。為了理解這個過程，我們將研究Julia系統的四個方面：動態語義、靜態編譯器內部函數的嵌入、過程間類型推斷，以及靜態子圖的提取。此外，我們還將研究這些特徵與巨集和生成的函數的互動，這些函數將與XLA編譯器相關。

如何將XLA嵌入到Julia IR

XLA嵌入

要編譯為XLA而不是LLVM，我們應用了上一節中概述的策略。實際上，我們可以重用大多數編譯器本身（特別是所有類型推斷和所有mid-level優化傳遞）。

讓我們先定義動態語義和靜態嵌入。

張量表示（Tensor representation）

由於其作為線性代數的教學和研究語言的傳統，Julia具有非常豐富的陣列抽象層次結構。Julia的標準庫陣列是可變的，並且在類型和維度上進行參數化。此外，StaticArrays.jl（Ferris＆Contributors，2018）包提供了在元素類型和形狀上進行參數化的不可變陣列。因此，成形的N維不可變張量的概念對Julia代碼來說並不陌生，並且大多數現有的通用代碼能夠毫無問題地處理它。

因此，我們通過定義一個runtime結構來嵌入XLA values。

Listing 1: XRTArray3的定義。

操作表示（Operation representation）

分離靜態和動態操作數

HLO操作數（HLO operands）分為靜態和動態操作數。假設我們有一個示例XLA操作'Foo'採用一個靜態操作數（例如一個整數）和兩個動態操作數。這個嵌入如下所示：

在這個示例中，“execute”函數實現在遠程設備上運行操作的動態語義。函數(hlo::HloFoo)(...) 語法表示調用運算符重載。因此，這意味著對HloFoo(1) 的調用將構造並返回一個callabale對象，當在兩個XRTArrays上調用時，它將使用靜態操作數'1'遠程執行'Foo'HLO操作，並且對應於兩個陣列的動態操作數。這種分離並不是絕對必要的，但確實有嵌入到Julia IR的有用特性，易於理解：

在Listing 2的示例中，我們將HLO操作數（包括靜態操作數）拚接到AST中。這產生了一個非常簡單的XLA映射（遍歷每個語句，從拚接指令規範獲取靜態操作數，從類型推斷獲得動態形狀並生成相應的XLA代碼）。

當然，我們通常不會手動拚接這些指令，但是手動拚接的示例說明了為什麽分離靜態操作數很有用，並說明了成功offload到XLA的條件。

如果經過所有相關的Julia級別優化之後，IR可以完全解除安裝：

Listing 2: 手動構建的XLA嵌入

滿足這些條件的IR可以簡單地轉換成XLA IR。

結果

本文描述的方法在很大程度上依賴於Julia中間端編譯器，以確定足夠精確的資訊，在程式的足夠大的子區域中分攤任何啟動開銷。

在本節中，我們證明了Julia編譯器確實足夠精確，使該方法適用於實際的程式。

VGG19 forward pass

圖1：在編譯到XLA之後，Metalhead.jl VGG19的forward pass 和backwards pass 生成的XLA指令摘要。

這裡顯示了未優化（在Julia前端之後）和優化的計數（在類似於CPU後端使用的XLA優化pipeline之後，但沒有HLO融合）。

VGG19 backward pass

為了獲得backwards pass，我們使用基於Zygote.jl編譯器的AD框架(Innes, 2018)。Zygote對Julia代碼進行操作，其輸出也是Julia函數（適合重新引入Zygote以獲得更高階導數，也適合編譯到TPU）。

示例如下：

結論

在這篇論文中，我們討論了如何將Julia代碼編譯為XLA IR，從而實現解除安裝到TPU設備。這裡描述的實現重新利用了現有Julia編譯器的重要部分，因此所有代碼不到1000行，但是仍然能夠編譯模型的forward和backward pass（及其融合，包括 training loop）到單個XLA內核，模型例如VGG19。

我們還演示了Julia的多重調度語義如何在這個轉換的規範中提供幫助。這項工作表明，不僅可以將用Julia編寫的多個ML模型編譯到TPU，而且可以編寫更通用的非ML Julia代碼（只要這些代碼也由線性代數操作控制）。我們希望這可以加速對非ML問題領域的探索，TPU可能對這些領域有用。

論文地址：

https://arxiv.org/pdf/1810.09868.pdf

新智元AI WORLD 2018

世界人工智能峰會全程回顧

新智元於9月20日在北京國家會議中心舉辦AI WORLD 2018世界人工智能峰會，邀請機器學習教父、CMU教授 Tom Mitchell，邁克思·泰格馬克，周志華，陶大程，陳怡然等AI領袖一起關注機器智能與人類命運。

全程回顧新智元 AI World 2018 世界人工智能峰會盛況：

愛奇藝

上午：https://www.iqiyi.com/v_19rr54cusk.html

下午：https://www.iqiyi.com/v_19rr54hels.html

新浪：http://video.sina.com.cn/l/p/1724373.html