如何在Julia程式設計中實現GPU加速

為了簡化操作，可以在 nextjournal 上註冊賬戶，點選「edit」即可直接執行文章中的簡單程式碼了。

註冊地址：https://nextjournal.com/signup

首先，什麼是 GPU？

GPU 是一種大型並行處理器，有幾千個並行處理單元。例如，本文使用的 Tesla k80，能提供 4992 個並行 CUDA 核。GPU 在頻率、延遲和硬體效能方面與 CPU 有很大的不同，但實際上 Tesla k80 有點類似於具有 4992 核的慢速 CPU。

能夠啟動的並行執行緒可以大幅提升速度，但也令使用 GPU 變得更困難。當使用這種未加處理的能量時，會出現以下缺點：

• GPU 是一種有專屬記憶體空間和不同架構的獨立硬體。因此，從 RAM 到 GPU 記憶體（VRAM，視訊記憶體）的傳輸時間很長。甚至在 GPU 上啟動核心（呼叫排程函式）也會帶來很大的延遲，對於 GPU 而言是 10us 左右，而對於 CPU 只有幾納秒。

• 在沒有高階封裝的情況下，建立核心會變得複雜。

• 低精度是預設值，高精度的計算可以很容易地消除所有效能增益。

• GPU 函式（核心）本質上是並行的，所以編寫 GPU 核心不比編寫並行 CPU 程式碼容易，而且硬體上的差異增加了一定的複雜性。

• 與上述情況相關的很多演算法都不能很好地遷移到 GPU 上。想要了解更多的細節，請看這篇博文：https://streamhpc.com/blog/2013-06-03/the application-areas-opencl- cuda-can- used/。

• 核心通常是用 C/ C++語言編寫的，但這並不是寫演算法的最好語言。

• CUDA 和 OpenCL 之間有差異，OpenCL 是編寫底層 GPU 程式碼的主要框架。雖然 CUDA 只支援英偉達硬體，OpenCL 支援所有硬體，但並不精細。要看個人需求進行選擇。

而Julia作為一種高階指令碼語言，允許在其中編寫核心和環境程式碼，同時可在大多數 GPU 硬體上執行！

GPUArrays

大多數高度並行的演算法都需要同時處理大量資料，以克服所有的多執行緒和延遲損耗。因此，大多數演算法都需要陣列來管理所有資料，這就需要一個好的 GPU 陣列庫作為關鍵的基礎。

GPUArrays.jl 是Julia為此提供的基礎。它實現了一個專門用於高度並行硬體的抽象陣列。它包含了設定 GPU、啟動JuliaGPU 函式、提供一些基本陣列演算法等所有必要功能。

抽象意味著它需要以 CuArrays 和 CLArrays 的形式實現。由於繼承了 GPUArrays 的所有功能，它們提供的介面完全相同。唯一的區別出現在分配陣列時，這會強制使用者決定這一陣列是存在於 CUDA 還是 OpenCL 裝置上。關於這一點的更多資訊，請參閱「記憶體」部分。

GPUArrays 有助於減少程式碼重複，因為它允許編寫獨立於硬體的 GPU 核心，這些核心可以通過 CuArrays 或 CLArrays 編譯到本地的 GPU 程式碼。因此，大多通用核心可以在從 GPUArrays 繼承的所有包之間共享。

選擇小貼士：CuArrays 只支援 Nvidia GPU，而 CLArrays 支援大多數可用的 GPU。CuArrays 比 CLArrays 更穩定，可以在Julia0.7 上使用。速度上兩者大同小異。我建議都試一試，看看哪種最有效。

本文中，我將選擇 CuArrays，因為本文是在Julia0.7 / 1.0 上編寫的，CLArrays 暫不支援。

效能

用一個簡單的互動式程式碼示例來快速說明：為了計算 julia 集合（曼德勃羅集合），我們必須要將計算轉移到 GPU 上。

using CuArrays, FileIO, Colors, GPUArrays, BenchmarkTools
using CuArrays: CuArray
"""
The function calculating the <mark data-type="technologies" data-id="2194d698-2e0a-4e6c-8f36-92d01507185d">Julia</mark> set
"""
function juliaset(z0, maxiter)
c = ComplexF32(-0.5, 0.75)
z = z0
for i in 1:maxiter
abs2(z) > 4f0 && return (i - 1) % UInt8
z = z * z + c
end
return maxiter % UInt8 # % is used to convert without overflow check
end
range = 100:50:2^12
cutimes, jltimes = Float64[], Float64[]
function run_bench(in, out)
# use dot syntax to apply `juliaset` to each elemt of q_converted 
# and write the output to result
out .= juliaset.(in, 16)
# all calls to the GPU are scheduled asynchronous, 
# so we need to synchronize
GPUArrays.synchronize(out)
end
# store a reference to the last results for plotting
last_jl, last_cu = nothing, nothing
for N in range
w, h = N, N
q = [ComplexF32(r, i) for i=1:-(2.0/w):-1, r=-1.5:(3.0/h):1.5]
for (times, Typ) in ((cutimes, CuArray), (jltimes, Array))
# convert to Array or CuArray - moving the calculation to CPU/GPU
q_converted = Typ(q)
result = Typ(zeros(UInt8, size(q)))
for i in 1:10 # 5 samples per size
# benchmarking macro, all variables need to be prefixed with $
t = Base.@elapsed begin
run_bench(q_converted, result)
end
global last_jl, last_cu # we're in local scope
if result isa CuArray
last_cu = result
else
last_jl = result
end
push!(times, t)
end
end
end

cu_jl = hcat(Array(last_cu), last_jl)
cmap = colormap("Blues", 16 + 1)
color_lookup(val, cmap) = cmap[val + 1]
save("results/juliaset.png", color_lookup.(cu_jl, (cmap,)))

using Plots; plotly()
x = repeat(range, inner = 10)
speedup = jltimes ./ cutimes
Plots.scatter(
log2.(x), [speedup, fill(1.0, length(speedup))], 
label = ["cuda" "cpu"], markersize = 2, markerstrokewidth = 0,
legend = :right, xlabel = "2^N", ylabel = "speedup"
)

對於大型陣列，通過將計算轉移到 GPU，可以穩定地將速度提高 60-80 倍。獲得此加速和將Julia陣列轉換為 GPUArray 一樣簡單。

有人可能認為 GPU 效能會受到像Julia這樣的動態語言影響，但Julia的 GPU 效能應該與 CUDA 或 OpenCL 的原始效能相當。Tim Besard 在整合 LLVM Nvidia 編譯流程方面做得很好，能夠實現與純 CUDA C 語言程式碼相同（有時甚至更好）的效能。他在部落格（https://devblogs.nvidia.com/gpu-computing-julia-programming-language/）中作了進一步解釋。CLArrays 方法有點不同，它直接從Julia生成 OpenCL C 程式碼，程式碼效能與 OpenCL C 相同！

為了更好地瞭解效能並與多執行緒 CPU 程式碼進行比對，我整理了一些基準：https://github.com/JuliaGPU/GPUBenchmarks.jl/blob/master/results/results.md

記憶體

GPU 具有自己的儲存空間，包括視訊記憶體（VRAM）、不同的快取記憶體和暫存器。無論做什麼，執行前都要先將Julia物件轉移到 GPU。並非Julia中的所有型別都可以在 GPU 上執行。

首先讓我們看一下Julia的型別：

struct Test # an immutable struct
# that only contains other immutable, which makes 
# isbitstype(Test) == true
x::Float32 
end

# the isbits property is important, since those types can be used
# without constraints on the GPU!
@assert isbitstype(Test) == true
x = (2, 2)
isa(x, Tuple{Int, Int}) # tuples are also immutable
mutable struct Test2 #-> mutable, isbits(Test2) == false
x::Float32
end
struct Test3
# contains a heap allocation/ reference, not isbits
x::Vector{Float32}
y::Test2 # Test2 is mutable and also heap allocated / a reference
end
Vector{Test} # <-An Array with isbits elements is contigious in memory
Vector{Test2} # <- An Array with mutable elements is basically an array of heap pointers. Since it just contains cpu heap pointers, it won't work on the GPU.

"Array{Test2,1}"

所有這些Julia型別在傳輸到 GPU 或在 GPU 上建立時表現不同。下表概述了預期結果：

建立位置描述物件是在 CPU 上建立的，然後轉移到 GPU 核心上，或者本身就由核心內部的 GPU 建立。該表顯示建立型別的例項是否可行，對於從 CPU 到 GPU 的轉移，該表還說明了物件是否能通過參照進行復制或傳遞。

垃圾收集

當使用 GPU 時，要注意 GPU 上沒有垃圾收集器（GC）。這不會造成太大影響，因為寫入 GPU 的高效能核心不應該建立任何 GC-跟蹤的記憶體作為起始。

在 GPU 上實現 GC 不無可能，但請記住，每個執行核心都是大規模並行的。在大約 1000 個 gpu 執行緒中的每一個建立和跟蹤大量堆記憶體就會馬上破壞效能增益，因此實現 GC 是得不償失的。

使用 GPUArrays 可以作為在核心中分配陣列的替代方法。GPUArray 建構函式將建立 GPU 緩衝區並將資料轉移到 VRAM。如果呼叫 Array(gpu_array)，陣列將被轉移回 RAM，變為普通的Julia陣列。這些 gpu 陣列的Julia操作由Julia的 GC 跟蹤，如果不再使用，GPU 記憶體將被釋放。

因此，只能在裝置上使用堆疊分配，並且只能被其他的預先分配的 GPU 緩衝區使用。由於轉移代價很高，因此在編寫 GPU 時，往往要儘可能重用和預分配。

GPUArray 建構函式

using CuArrays, LinearAlgebra

# GPU Arrays can be constructed from all <mark data-type="technologies" data-id="2194d698-2e0a-4e6c-8f36-92d01507185d">Julia</mark> arrays containing isbits types!
A1D = cu([1, 2, 3]) # cl for CLArrays
A1D = fill(CuArray{Int}, 0, (100,)) # CLArray for CLArrays
# Float32 array - Float32 is usually preferred and can be up to 30x faster on most GPUs than Float64
diagonal_matrix = CuArray{Float32}(I, 100, 100)
filled = fill(CuArray, 77f0, (4, 4, 4)) # 3D array filled with Float32 77
randy = rand(CuArray, Float32, 42, 42) # random numbers generated on the GPU
# The array constructor also accepts isbits iterators with a known size
# Note, that since you can also pass isbits types to a gpu kernel directly, in most cases you won't need to materialize them as an gpu array
from_iter = CuArray(1:10)
# let's create a point type to further illustrate what can be done:
struct Point
x::Float32
y::Float32
end
Base.convert(::Type{Point}, x::NTuple{2, Any}) = Point(x[1], x[2])
# because we defined the above convert from a tuple to a point
# [Point(2, 2)] can be written as Point[(2,2)] since all array 
# elements will get converted to Point
custom_types = cu(Point[(1, 2), (4, 3), (2, 2)])
typeof(custom_types)

"CuArray{point,1}"

陣列操作

我們已經定義了許多操作。最重要的是，GPUArrays 支援Julia的融合點廣播表示法（fusing dot broadcasting notation）。此表示法允許你將函式應用於陣列的每個元素，並使用 f 的返回值建立新陣列。此功能通常稱為對映（map）。broadcast 指的是形狀各異的陣列被 broadcast 成相同形狀。

工作原理如下：

x = zeros(4, 4) # 4x4 array of zeros
y = zeros(4) # 4 element array
z = 2 # a scalar
# y's 1st dimension gets repeated for the 2nd dimension in x
# and the scalar z get's repeated for all dimensions
# the below is equal to `broadcast(+, broadcast(+, xx, y), z)`
x .+ y .+ z

發生「融合」是因為Julia編譯器會重寫該表示式為一個傳遞呼叫樹的 lazy broadcast 呼叫，然後可以在迴圈遍歷陣列之前將整個呼叫樹融合到一個函式中。

如果你想要更詳細的瞭解 broadcast，可以看該指南：julia.guide/broadcasting。

這意味著在不分配堆記憶體（僅建立 isbits 型別）的情況下執行的任何Julia函式，都可以應用於 GPUArray 的每個元素，並且多點呼叫會融合到一個核心呼叫中。由於核心呼叫會有很大延遲，所以這種融合是一個非常重要的優化。

using CuArrays
A = cu([1, 2, 3])
B = cu([1, 2, 3])
C = rand(CuArray, Float32, 3)
result = A .+ B .- C
test(a::T) where T = a * convert(T, 2) # convert to same type as `a`

# inplace broadcast, writes directly into `result`
result .= test.(A) # custom function work

# The cool thing is that this composes well with custom types and custom functions.
# Let's go back to our Point type and define addition for it
Base.:(+)(p1::Point, p2::Point) = Point(p1.x + p2.x, p1.y + p2.y)

# now this works:
custom_types = cu(Point[(1, 2), (4, 3), (2, 2)])

# This particular example also shows the power of broadcasting: 
# Non array types are broadcasted and repeated for the whole length
result = custom_types .+ Ref(Point(2, 2))

# So the above is equal to (minus all the allocations):
# this allocates a new array on the gpu, which we can avoid with the above broadcast
broadcasted = fill(CuArray, Point(2, 2), (3,))

result == custom_types .+ broadcasted

true

GPUArrays 支援更多操作：

• 實現 GPU 陣列轉換為 CPU 陣列和複製

• 多維索引和切片 (xs[1:2, 5, :])

• permutedims

• 串聯 (vcat(x, y), cat(3, xs, ys, zs))

• 對映，融合 broadcast(zs .= xs.^2 .+ ys .* 2)

• 填充 (CuArray, 0f0, dims)，填充 (gpu_array, 0)

• 減小尺寸 (reduce(+, xs, dims = 3), sum(x -> x^2, xs, dims = 1)

• 縮減為標量 (reduce(*, xs), sum(xs), prod(xs))

• 各種 BLAS 操作 (matrix*matrix, matrix*vector)

• FFT，使用與 julia 的 FFT 相同的 API

GPUArrays 實際應用

讓我們直接看一些很酷的例項。

GPU 加速煙霧模擬器是由 GPUArrays + CLArrays 建立的，可在 GPU 或 CPU 上執行，GPU 版本速度提升 15 倍：

還有更多的例子，包括求微分方程、FEM 模擬和求解偏微分方程。

演示地址：https://juliagpu.github.io/GPUShowcases.jl/latest/index.html

讓我們通過一個簡單的機器學習示例，看看如何使用 GPUArrays：

using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated, partition
using CuArrays

# Classify MNIST digits with a convolutional network

imgs = MNIST.images()

labels = onehotbatch(MNIST.labels(), 0:9)

# Partition into batches of size 1,000
train = [(cat(float.(imgs[i])..., dims = 4), labels[:,i])
for i in partition(1:60_000, 1000)]

use_gpu = true # helper to easily switch between gpu/cpu

todevice(x) = use_gpu ? gpu(x) : x

train = todevice.(train)

# Prepare test set (first 1,000 images)
tX = cat(float.(MNIST.images(:test)[1:1000])..., dims = 4) |> todevice
tY = onehotbatch(MNIST.labels(:test)[1:1000], 0:9) |> todevice

m = Chain(
Conv((2,2), 1=>16, relu),
x -> maxpool(x, (2,2)),
Conv((2,2), 16=>8, relu),
x -> maxpool(x, (2,2)),
x -> reshape(x, :, size(x, 4)),
Dense(288, 10), softmax) |> todevice

m(train[1][1])

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(onecold(m(x)) .== onecold(y))

evalcb = throttle(() -> @show(accuracy(tX, tY)), 10)
opt = ADAM(Flux.params(m));

# train
for i = 1:10
Flux.train!(loss, train, opt, cb = evalcb)
end

using Colors, FileIO, ImageShow
N = 22
img = tX[:, :, 1:1, N:N]
println("Predicted: ", Flux.onecold(m(img)) .- 1)
Gray.(collect(tX[:, :, 1, N]))

只需將陣列轉換為 GPUArrays（使用 gpu(array），就可以將整個計算移動到 GPU 並獲得可觀的速度提升。這要歸功於Julia複雜的 AbstractArray 基礎架構，使 GPUArray 可以無縫整合。隨後，如果省略轉換為 GPUArray 這一步，程式碼會按普通的Julia陣列處理，但仍在 CPU 上執行。可以嘗試將 use_gpu = true 改為 use_gpu = false，重新執行初始化和訓練單元格。對比 GPU 和 CPU，CPU 執行時間為 975 秒，GPU 執行時間為 29 秒，速度提升約 33 倍。

另一個優勢是為了有效地支援神經網路的反向傳播，GPUArrays 無需明確地實現自動微分。這是因為Julia的自動微分庫適用於任意函式，並存有可在 GPU 上高效執行的程式碼。這樣即可利用最少的開發人員就能在 GPU 上實現 Flux，並使 Flux GPU 能夠高效實現使用者定義的功能。這種開箱即用的 GPUArrays + Flux 不需要協調，這是Julia的一大特點，詳細解釋如下：為什麼 Numba 和 Cython 不能代替Julia（http://www.stochasticlifestyle.com/why）。

編寫 GPU 核心

一般情況，只使用 GPUArrays 的通用抽象陣列介面即可，而不需要編寫任何 GPU 核心。但是有些時候，可能需要在 GPU 上實現一個無法通過一般陣列演算法組合表示的演算法。

好訊息是，GPUArrays 通過分層法消除了大量工作，可以實現從高階程式碼開始，編寫類似於大多數 OpenCL / CUDA 示例的低階核心。同時可以在 OpenCL 或 CUDA 裝置上執行核心，從而提取出這些框架中的所有差異。

實現上述功能的函式名為 gpu_call。呼叫語句為 gpu_call(kernel, A::GPUArray, args)，在 GPU 上使用引數(state, args...) 呼叫 kernel。State 是一個用於實現獲取執行緒索引等功能的後端特定物件。GPUArray 需要作為第二個引數傳遞，以分配到正確的後端並提供啟動引數的預設值。

讓我們使用 gpu_call 來實現一個簡單的對映核心：

using GPUArrays, CuArrays
# Overloading the <mark data-type="technologies" data-id="2194d698-2e0a-4e6c-8f36-92d01507185d">Julia</mark> Base map! function for GPUArrays
function Base.map!(f::Function, A::GPUArray, B::GPUArray)
# our function that will run on the gpu
function kernel(state, f, A, B)
# If launch <mark data-type="technologies" data-id="2e982b73-88e2-41e8-a430-f7ae5a9af4bf">parameter</mark>s aren't specified, linear_index gets the index
# into the Array passed as second argument to gpu_call (`A`)
i = linear_index(state)
if i <= length(A)
@inbounds A[i] = f(B[i])
end
return
end
# call kernel on the gpu
gpu_call(kernel, A, (f, A, B))
end

簡單來說，這將在 GPU 上並行呼叫 julia 函式 kernel length(A) 次。kernel 的每個並行呼叫都有一個執行緒索引，可以利用它索引到陣列 A 和 B。如果計算索引時沒有使用 linear_index，就需要確保沒有多個執行緒讀取和寫入相同的陣列位置。因此，如果在純Julia中使用執行緒編寫，可等效如下：

using BenchmarkTools
function threadded_map!(f::Function, A::Array, B::Array)
Threads.@threads for i in 1:length(A)
A[i] = f(B[i])
end
A
end
x, y = rand(10^7), rand(10^7)
kernel(y) = (y / 33f0) * (732.f0/y)
# on the cpu without threads:
single_t = @belapsed map!($kernel, $x, $y)

# "on the CPU with 4 threads (2 real cores):
thread_t = @belapsed threadded_map!($kernel, $x, $y)

# on the GPU:
xgpu, ygpu = cu(x), cu(y)
gpu_t = @belapsed begin
map!($kernel, $xgpu, $ygpu)
GPUArrays.synchronize($xgpu)
end
times = [single_t, thread_t, gpu_t]
speedup = maximum(times) ./ times
println("speedup: $speedup")
bar(["1 core", "2 cores", "gpu"], speedup, legend = false, fillcolor = :grey, ylabel = "speedup")

由於該函式未實現過多內容，也得不到更多的擴充套件，但執行緒化和 GPU 版本仍然有一個很好的加速。

GPU 與執行緒示例相比，能顯示更復雜的內容，因為硬體執行緒是以執行緒塊的形式分佈的，gpu_call 是從簡單版本中提取出來的，但它也可以用於更復雜的啟動配置：

using CuArrays

threads = (2, 2)
blocks = (2, 2)
T = fill(CuArray, (0, 0), (4, 4))
B = fill(CuArray, (0, 0), (4, 4))
gpu_call(T, (B, T), (blocks, threads)) do state, A, B
# those names pretty much refer to the cuda names
b = (blockidx_x(state), blockidx_y(state))
bdim = (blockdim_x(state), blockdim_y(state))
t = (threadidx_x(state), threadidx_y(state))
idx = (bdim .* (b .- 1)) .+ t
A[idx...] = b
B[idx...] = t
return
end
println("Threads index: \n", T)
println("Block index: \n", B)

上面的示例中啟動配置的迭代順序更復雜。確定合適的迭代+啟動配置對於實現最優 GPU 效能至關重要。很多關於 CUDA 和 OpenCL 的 GPU 教程都非常詳細地解釋了這一點，在Julia中程式設計 GPU 時這些原理是相通的。

結論

Julia為高效能的世界帶來了可組合的高階程式設計。現在是時候為 GPU 做同樣的事了。

希望Julia能降低人們在 GPU 程式設計的門檻，我們可以為開源 GPU 計算開發可擴充套件的平臺。第一個成功案例是通過Julia軟體包實現自動微分解決方案，這些軟體包甚至都不是為 GPU 編寫的，因此可以相信Julia在 GPU 計算領域的擴充套件性和通用設計中一定會大放異彩。

原文連結：https://nextjournal.com/sdanisch/julia-gpu-programming