Numba, Julia的GPU編程(也許會再寫點XLA)

來自專欄 Half Integer12 人贊了文章

本來想著後天講點這個，不過感覺只有半個小時，我還要講Julia的OO（畢竟這個現階段應該更實用一點）其實講不了多少。就寫個文章簡單扯一扯

會了解這個主要還是因為之前有考慮給我們最近寫的量子線路的模擬器Yao.jl上GPU。不過我還沒有了解太深入其實，所以還請補（tou）充（gao）和糾（ping）正（lun）

為什麼要扯tensorflow？除了蹭機器學習的熱點以外，還有是因為Julia的GPU實現和tensorflow的JIT後端語言XLA有一些相似之處，起初知道這個是看到了有人在hacker news上這麼說：

This is in part because of the work by Google on the NVPTX LLVM back-end.

Im one of the maintainers at Google of the LLVM NVPTX backend. Happy to answer questions about it.

As background, Nvidias CUDA ("CUDA C++?") compiler, nvcc, uses a fork of LLVM as its backend. Clang can also compile CUDA code, using regular upstream LLVM as its backend. The relevant backend in LLVM was originally contributed by nvidia, but these days the team Im on at Google is the main contributor.

I dont know much (okay, anything) about Julia except what I read in this blog post, but the dynamic specialization looks a lot like XLA, a JIT backend for TensorFlow that I work on. So thats cool; Im happy to see this work.

Full debug information is not supported by the LLVM NVPTX back-end yet, so cuda-gdb will not work yet.

Wed love help with this.

Bounds-checked arrays are not supported yet, due to a bug [1] in the NVIDIA PTX compiler. [0]

We ran into what appears to be the same issue [2] about a year and a half ago. nvidia is well aware of the issue, but I dont expect a fix except by upgrading to Volta hardware.

[0] https://julialang.org/blog/2017/03/cudanative [1] https://github.com/JuliaGPU/CUDAnative.jl/issues/4 [2] https://bugs.llvm.org/show_bug.cgi?id=27738

然後我們就先看看Julia的編譯過程吧，這個我稍微熟悉一點，下面的主要內容都來自CUDAnative作者的論文（也參考了趙岳華的slide）:

Unleashing Julia on GPUs?

arxiv.org

本身Julia語言的編譯器大約由兩層IR+LLVM IR的代碼生成構成，所以實際上除了LLVM自己的東西，Julia編譯器裡面也有一些自帶的優化。

1. 首先編譯器讀入源代碼經過lexer tokenize，然後parse成一個表層的AST，這個AST可以直接在Julia裡面通過 Expr 類型，QuoteNode類型和Symbol類型獲得（實際上token基本上就可以認為是Symbol），這個Expr長的和S-Expr差不多，這一部分主要是一個用FemoteLisp寫的解析器完成的。這一層的IR可以通過 引用 的語法來獲得，並且修改（通過宏）
2. 接下來會展開Julia語法裡面的宏，獲得更簡單的表示，然後交給高級優化器優化，獲得下層的IR，這一層將會進行類型推導（type inference）可以通過 generated function 來修改。
3. 然後經過低級優化器進入codegen產生LLVM IR之後會運行到LLVM的CPU後端上
4. 如果在表層AST被展開時發現有一段代碼被標記了要在GPU上運行（通過cuda宏），那麼這一部分就會進入CUDAnative的這些編譯器extension

CUDAnative實際上就是通過修改（作為一個編譯器的擴展）上面的這些IR來實現將Julia源代碼編譯到PTX的目標。於是你可以直接復用大量的Julia源代碼到GPU上去：

using CUDAnative, CUDAdrvfunction kernel_vadd(a, b, c) i = (blockIdx().x-1) * blockDim().x + threadIdx().x c[i] = a[i] + b[i] return nothingenda = round.(rand(Float32, (3, 4)) * 100)b = round.(rand(Float32, (3, 4)) * 100)d_a = CuArray(a)d_b = CuArray(b)d_c = similar(d_a)@cuda threads=12 kernel_vadd(d_a, d_b, d_c) # 注意這個是0.7+里才能用的，0.6必須用Tuple指明 block 和dim

值得注意的是這個編譯器擴展本身也是Julia寫的：

rogerluo@amax:~/.julia/v0.6/CUDAnative$ cloc .

74 text files.

74 unique files.

614 files ignored.

CLOC -- Count Lines of Code v 1.60 T=0.25 s (192.9 files/s, 19896.3 lines/s)

-------------------------------------------------------------------------------

Language files blank comment code

-------------------------------------------------------------------------------

Lisp 43 1095 0 3695

C 1 24 6 71

make 2 9 0 23

YAML 2 5 0 22

-------------------------------------------------------------------------------

SUM: 48 1133 6 3811

-------------------------------------------------------------------------------

那個Lisp就是Julia，cloc經常能把Julia識別為Lisp...

然後嘛，我們說Python的numba也能寫CUDA啊，有啥區別呢？主要是編譯過程不一樣

這個就表現為只有少量numba支持的Python語法和類型才能被用來產生PTX。Numba自己實現了一套自己的類型系統和IR然後直接將Python編譯出來的位元組碼拿去修改。

放幾張論文里的圖

就代碼量上來說Julia確實比CUDA C要少不少，但是性能上其實是有一定損失的，作者報告說大約會比CUDA C損失0.5%左右

理論上講CUDAnative和其一系列衍生的包：GPUArrays，CuArrays能夠支持所有用戶自定義類型和函數。不過就實際體驗來說還是別真的來在Julia裡面寫CUDA，其實還有一堆bug，17年底才基本寫好???♂?，文檔都不全的，如果你不打算想我一樣看他實現的話，還是得等上幾個月或者一年再來用。

不過論調庫的話，調個cublas和cudnn還是基本OK的，然而這個除了一堆語法糖和自定義的函數支持以外應該和你用PyTorch在Python里的實際使用感受差不多（像ForwardDiff這些自動微分庫一般是支持的StridedArray，所以不管你什麼Array都可以自動微分）。