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

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

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

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

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

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

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

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

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

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

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

谷歌TPU和XLA編譯器

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

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

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

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

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

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

Julia編譯器的工作原理

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

Julia編譯器需要將語言的動(dòng)態(tài)語義與LLVM表示的靜態(tài)語義之間聯(lián)系起來。為了理解這個(gè)過程，我們將研究Julia系統(tǒng)的四個(gè)方面：動(dò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í)際上，我們可以重用大多數(shù)編譯器本身（特別是所有類型推斷和所有mid-level優(yōu)化傳遞）。

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

張量表示（Tensor representation）

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

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

Listing 1: XRTArray3的定義。

操作表示（Operation representation）

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

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

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

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

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

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

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

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

結(jié)果

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

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

VGG19 forward pass

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

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

VGG19 backward pass

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

示例如下：

結(jié)論

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

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