將Julia代碼直接部署到谷歌Cloud TPU，讓程序運行更快的官方指南來了！Julia和TPU的結(jié)合意味著快速、易于表達的ML計算！”

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

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

隨著用戶越來越多，圍繞Julia的開發(fā)工具、技術(shù)、教程等也愈加豐富。昨天，Julia開發(fā)人員Keno Fischer和Elliot Saba發(fā)表了一篇新論文AutomaticFullCompilationof JuliaProgramsandMLModelstoCloudTPUs，介紹如何將Julia代碼直接部署到Google Cloud TPU，讓程序運行更快。

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

谷歌的Cloud TPU是一種用于機器學(xué)習(xí)工作負載的很強大的新硬件架構(gòu)。近年來，Cloud TPU為谷歌的許多里程碑式的機器學(xué)習(xí)成就提供了動力。

谷歌現(xiàn)在已經(jīng)在他們的云平臺上開放提供一般用途的TPU，并且最近已經(jīng)進一步開放，允許非TensorFlow前端使用。

這篇論文描述了通過這個新的API和Google XLA編譯器，將Julia程序的適當(dāng)部分卸載（offload）到TPU的方法和實現(xiàn)。

這一方法能夠?qū)⒈硎緸镴ulia程序的VGG19模型的前向傳遞（forward pass）完全融合到單個TPU可執(zhí)行文件中，以便卸載到設(shè)備。該方法也很好地與Julia代碼上現(xiàn)有的基于編譯器的自動微分技術(shù)相結(jié)合，因此我們也能夠自動獲得VGG19的反向傳遞并類似地將其卸載到TPU。

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

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

谷歌TPU和XLA編譯器

2017年，谷歌宣布他們將通過云服務(wù)向公眾提供其專有的張量處理單元（TPU）機器學(xué)習(xí)加速器。最初，TPU的使用僅限于使用谷歌的TensorFlow機器學(xué)習(xí)框架編寫的應(yīng)用程序。幸運的是，2018年9月，Google通過較低級別的XLA（Accelerated Linear Algebra）編譯器的IR開放了對TPU的訪問權(quán)限。該IR是通用的，是用于表示線性代數(shù)原語的任意計算的優(yōu)化編譯器，因此為非Tensorflow用戶以及非機器學(xué)習(xí)工作負載的TPU目標(biāo)提供了良好的基礎(chǔ)。

XLA（加速線性代數(shù)）是谷歌的一個部分開源編譯器項目。它具有豐富的輸入IR，用于指定多線性代數(shù)計算，并為CPU，GPU和TPU提供后端代碼生成功能。XLA的輸入IR（稱為HLO高級優(yōu)化IR）在基本數(shù)據(jù)類型或其元組（但沒有元組數(shù)組）的任意維數(shù)組上運行。HLO操作包括基本算術(shù)運算、特殊函數(shù)、廣義線性代數(shù)運算、高級數(shù)組運算以及用于分布式計算的原語。XLA可以執(zhí)行輸入程序的語義簡化，以及執(zhí)行整個程序的內(nèi)存調(diào)度，以便有效地使用和重用可用內(nèi)存（這是大型機器學(xué)習(xí)模型的一個非常重要的考慮因素）。

每個HLO操作都有兩種操作數(shù)：

靜態(tài)操作數(shù)，它的值必須在編譯時可用并配置操作。

動態(tài)操作數(shù)，由上述張量組成。

這篇論文介紹了使用這個接口將常規(guī)的Julia代碼編譯帶TPU的初步工作。這一方法不依賴跟蹤，而是利用Julia的靜態(tài)分析和編譯功能來編譯完整的程序，包括對設(shè)備的任何控制flow。特別是，我們的方法允許用戶在編寫模型時充分利用Julia語言的完整表現(xiàn)力，能夠編譯使用Flux機器學(xué)習(xí)框架編寫的完整機器學(xué)習(xí)模型，將前向和后向模型傳遞以及訓(xùn)練loop融合到單個可執(zhí)行文件，并將其卸載到TPU。

Julia編譯器的工作原理

為了理解如何將Julia代碼編譯為XLA代碼，了解常規(guī)Julia編譯器的工作原理是有益的。Julia在語義上是一種非常動態(tài)的語言。但是，在標(biāo)準(zhǔn)配置中，Julia的最終后端編譯器是LLVM（Lattner＆Adve，2004），它是一個靜態(tài)編譯器后端。

Julia編譯器需要將語言的動態(tài)語義與LLVM表示的靜態(tài)語義之間聯(lián)系起來。為了理解這個過程，我們將研究Julia系統(tǒng)的四個方面：動態(tài)語義、靜態(tài)編譯器內(nèi)部函數(shù)的嵌入、過程間類型推斷，以及靜態(tài)子圖的提取。此外，我們還將研究這些特征與宏和生成的函數(shù)的交互，這些函數(shù)將與XLA編譯器相關(guān)。

如何將XLA嵌入到Julia IR

XLA嵌入

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

讓我們先定義動態(tài)語義和靜態(tài)嵌入。

張量表示（Tensor representation）

由于其作為線性代數(shù)的教學(xué)和研究語言的傳統(tǒng)，Julia具有非常豐富的數(shù)組抽象層次結(jié)構(gòu)。Julia的標(biāo)準(zhǔn)庫數(shù)組是可變的，并且在類型和維度上進行參數(shù)化。此外，StaticArrays.jl（Ferris＆Contributors，2018）包提供了在元素類型和形狀上進行參數(shù)化的不可變數(shù)組。因此，成形的N維不可變張量的概念對Julia代碼來說并不陌生，并且大多數(shù)現(xiàn)有的通用代碼能夠毫無問題地處理它。

因此，我們通過定義一個runtime結(jié)構(gòu)來嵌入XLA values。

Listing 1: XRTArray3的定義。

操作表示（Operation representation）

分離靜態(tài)和動態(tài)操作數(shù)

HLO操作數(shù)（HLO operands）分為靜態(tài)和動態(tài)操作數(shù)。假設(shè)我們有一個示例XLA操作'Foo'采用一個靜態(tài)操作數(shù)（例如一個整數(shù)）和兩個動態(tài)操作數(shù)。這個嵌入如下所示：

在這個示例中，“execute”函數(shù)實現(xiàn)在遠程設(shè)備上運行操作的動態(tài)語義。函數(shù)(hlo::HloFoo)(...) 語法表示調(diào)用運算符重載。因此，這意味著對HloFoo(1) 的調(diào)用將構(gòu)造并返回一個callabale對象，當(dāng)在兩個XRTArrays上調(diào)用時，它將使用靜態(tài)操作數(shù)'1'遠程執(zhí)行'Foo'HLO操作，并且對應(yīng)于兩個數(shù)組的動態(tài)操作數(shù)。這種分離并不是絕對必要的，但確實有嵌入到Julia IR的有用特性，易于理解：

在Listing 2的示例中，我們將HLO操作數(shù)（包括靜態(tài)操作數(shù)）拼接到AST中。這產(chǎn)生了一個非常簡單的XLA映射（遍歷每個語句，從拼接指令規(guī)范獲取靜態(tài)操作數(shù)，從類型推斷獲得動態(tài)形狀并生成相應(yīng)的XLA代碼）。

當(dāng)然，我們通常不會手動拼接這些指令，但是手動拼接的示例說明了為什么分離靜態(tài)操作數(shù)很有用，并說明了成功offload到XLA的條件。

如果經(jīng)過所有相關(guān)的Julia級別優(yōu)化之后，IR可以完全卸載：

Listing 2: 手動構(gòu)建的XLA嵌入

滿足這些條件的IR可以簡單地轉(zhuǎn)換成XLA IR。

結(jié)果

本文描述的方法在很大程度上依賴于Julia中間端編譯器，以確定足夠精確的信息，在程序的足夠大的子區(qū)域中分攤?cè)魏螁娱_銷。

在本節(jié)中，我們證明了Julia編譯器確實足夠精確，使該方法適用于實際的程序。

VGG19 forward pass

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

這里顯示了未優(yōu)化（在Julia前端之后）和優(yōu)化的計數(shù)（在類似于CPU后端使用的XLA優(yōu)化pipeline之后，但沒有HLO融合）。

VGG19 backward pass

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

示例如下：

結(jié)論

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

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