はじめに

前回の記事でMetal.jl を用いたGPGPUの現状について紹介しました。
Apple SiliconのGPUを使えるAPI としては、Metalの他にOpenCLがあります。

Julia にはOpenCLのバインディングであるOpenCL.jlもあり、これを利用してもGPU計算が可能なのでここで紹介します。

さらに、Array Abstractionに関しては、利便性の高いHigh-Level APIであるArrayFireのバインディングのバックエンドとしてOpenCL を利用することでカーネルを書かずにかなり最適化された行列計算が可能になっているので併せて紹介します。

筆者の環境は

MacBook Pro (14-inch, 2021)
Chip: Apple M1 Pro (14 Core GPU)
OS: Monterey 12.4
Julia v1.8.0-beta3

です。

Array Abstraction について

julia言語においては"Array Abstraction" として行列演算の並列化をユーザーに意識させずに可能にしようという考え方があります。これにより、AbstractArray型を引数にする関数を使っている限りにおいてとりあえず並列化が有効になった各種Array(CuArray etc...) をぶち込むだけで並列化が(しかもそれなりに最適化されて)有効になるという仕組みをとっています。
(徹底した型安定!)

プロセス並列化(DistributedArrays.jl)
スレッド並列化(Strided.jl)
GPU並列化(GPUArrays.jl)

そのGPU版がGPUArrays.jl です。実際にはこれを書くバックエンドとなる言語ごとに実装した CUDA.jlやAMDGPU.jl などのパッケージに多くの場合内包されています。

ただし、実際にどこまで計算に使うことができるかはパッケージによります。

CUDAの場合

CUDA.jlがjuliaの中では最も頻繁に更新されており、とんでもないことに、BLAS のLevel-3 までは全て実装されているのでBLAS にある計算は全て可能です。

直接実装はないため、左方除算\ に対応する計算はそのままではできません。getrfでLU分解してからgetrsで可能なようです。(cuSOLVER)

さらに、FFTまで実装されており、メモリの許す限り大きなベクトルも爆速でFFTしてくれます。(超便利)

OpenCLの場合

本題のOpenCL.jlですが、2年前に開発が途絶えています。
ArrayAbstractionのためのパッケージであるCLArrays.jl も残念ながら途絶えており、Julia v1.0 以降では利用できません。

それもある意味で当然で、これまではデバイスのベンダごとに CUDA.jl, AMDGPU.jl, oneAPI.jl が実装されているので基本的にJulia が動く環境が網羅できていました。Apple Silicon が過渡期であるからアクセスしにくいにすぎません。

いますぐ、Julia + Apple Silicon でGPGPUしたい!

こう思っている人は少なくないのではないでしょうか。そんな方々に朗報で、実は(前項のMetal.jlのようにカーネルを書かず)にできる、というのが本稿の動機になります。

本題

というわけで、OpenCL を利用してM1 GPU を使っていきましょう

インストール

pkg から可能です

pkg> add OpenCL

テスト

test.jl

using OpenCL
using BenchmarkTools

Mdim = 2^10
Pdim = 2^10
Ndim = 2^10

A = rand(Float32, Mdim, Pdim)
B = rand(Float32, Pdim, Ndim)
C = zeros(Float32, Mdim, Ndim)
C_cpu = similar(C)
# set up OpenCL

function bench_gpu!(A,B,C)
    M = size(A,1)
    P = size(A,2)
    N = size(B,2)
    ctx = cl.create_some_context()
    queue = cl.CmdQueue(ctx, :profile)
    # create OpenCL Buffers
    a_buff = cl.Buffer(Float32, ctx, (:r,:copy), hostbuf= vec(A))
    b_buff = cl.Buffer(Float32, ctx, (:r,:copy), hostbuf= vec(B))
    c_buff = cl.Buffer(Float32, ctx, :w, length(C))

    matmul_kernel = "
    __kernel void matmul(const int M, const int P, const int N,
                        const __global float* A,
                        const __global float* B,
                        __global float* C) {
        const int i = get_global_id(0);
        const int k = get_global_id(1);
        float tmp = 0.0f;
        for (int j=0; j<P; j++) {
            tmp += A[i + j*M] * B[j + k*P];
        }
        C[i + k*M] = tmp;
    }
    "

    prog = cl.Program(ctx, source = matmul_kernel) |> cl.build!
    kernel = cl.Kernel(prog, "matmul")
    run_time_vec = []
    ## warmup
    run_time = 0.0
    for i = 1:10
        evt = queue(kernel, (M, N), nothing, Int32(M), Int32(P), Int32(N),a_buff , b_buff , c_buff)
        run_time = evt[:profile_duration] ##ns
        run_time_vec = [run_time_vec; run_time]
    end
    GFLOPS = 2 * M * P * N /(minimum(run_time_vec))
    println("Peak Performance is $GFLOPS GFLOPS")

    evt = queue(kernel, (M, N), nothing, Int32(M), Int32(P), Int32(N),a_buff , b_buff , c_buff)
    cl.copy!(queue, C, c_buff)
    nothing
end

function bench_cpu!(A,B,C_cpu)
    C_cpu = A * B
    b = @belapsed $A*$B
    GFLOPS = 2 * length(A) * size(B,2) / b /1e9
    println("Peak Performance is $GFLOPS GFLOPS")
    nothing
end

筆者の環境で実行するとこのような感じ

julia> bench_gpu!(A,B,C)
Peak Performance is 10878.677872170128 GFLOPS

julia> bench_cpu!(A,B,C_cpu)
Peak Performance is 285.9118157369192 GFLOPS

julia> C ≈ A * B
true

ちょっと計測はうまくいっていない(理論性能の4.6TFLOPS を超えている!)ようで、表示上は速く見えますが実際はCPUと変わりません。

julia> @time bench_gpu!(A,B,C)
Peak Performance is 10580.50928726979 GFLOPS
  0.122489 seconds (1.22 k allocations: 71.594 KiB)
julia>  @time A * B;
  0.010347 seconds (2 allocations: 4.000 MiB)

bench_gpu!(A,B,C) では合計11回かけていることを加味してもCPUと同じくらいです。
ここから最適化.....
だったらMetal.jl 使えよ! となると思います。

ここから本番(ArrayFire.jl)

ArrayFire とは、パフォーマンス重視の並列計算用の高水準API で、Julia のAbstract Array の考え方と全く同様に、配列ベースで気軽に並列計算できるようにするためのツールです。
Julia と同様に複数のバックエンド(CUDA, OpenCL, CPU) から選択してバックではその言語で動作し、ユーザーはコードをそのまま異なるプラットフォームに移植できます。

このパッケージも同様にメンテナンスされていません。しかし、現在のバージョンでも動作することがわかりました。(少し操作が必要です。)

インストール

ArrayFire.jl をクローンします。

$ git clone https://github.com/JuliaGPU/ArrayFire.jl && cd ArrayFire.jl

Project.toml の記述が古いです。実際にはこれらのパッケージを新しくしても動作するので書き換えてしまっても問題は生じていません・

Project.toml

name = "ArrayFire"
uuid = "b19378d9-d87a-599a-927f-45f220a2c452"
authors = ["Julia Computing"]
version = "1.0.7"

[deps]
DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[compat]
DSP = "0.5.2,0.6,0.7"
FFTW = "0.2.4,1"
SpecialFunctions = "0.7,0.8,0.9,0.10,1,2"
julia = "1"

[compat] にDSP と SpecialFunctions のバージョンを追加しただけです。この状態でビルドしましょう

pkg> build
julia> using ArrayFire

メインの環境に直接インストールしてしまいたい場合は Git をForkしてそのレポジトリで書き換えるのがいいと思います。
自分はProject.toml のみ書き換えたFork を作成して直でいれています。

pkg> add https://github.com/IvanKuznetsoff/ArrayFire.jl
julia> using ArrayFire
ArrayFire v3.8.1 (OpenCL, 64-bit Mac OSX, build default)
[0] APPLE: Apple M1 Max, 49152 MB

何も環境変数をいじっていない状態だとOpenCL バックエンドが有効化されています。
Apple Silicon なのでシングルGPUとしては異常な 50GB近くのメモリを何も考えずに利用できます!

テスト

using ArrayFire
ArrayFire v3.8.1 (OpenCL, 64-bit Mac OSX, build default)
[0] APPLE: Apple M1 Max, 49152 MB
julia> A = rand(Float32, 2^10, 2^10); B = rand(Float32, 2^10, 2^10);
julia> AF_A = AFArray(A); AF_B = AFArray(B);
julia> @time AF_A * AF_B;
  0.002057 seconds (1 allocation: 16 bytes)
julia> @time A * B;
  0.013057 seconds (2 allocations: 4.000 MiB)

大体6倍早くなりました。なかなかいい感じ。

嬉しいのは batched gemm がデフォルトで使えるところです。
これは計算としては
$$
C_{ikn} = \sum_j A_{ijn}B_{jkn}
$$
を行うもので、例えばGradient Descent のヤコビアンを並列で走らせようと考える時、Naive にスレッド並列化をしてしまうとメモリのアクセスが離散的になりますが、勾配をさらに一次元大きい行列にしてしまえばメモリアクセスが連続的になって早くなるというメリットがあり、(よく?) 機械学習で出てくる計算だと思います。

julia> using NNlib
julia> A = rand(Float32, 2^10, 2^10, 2^10); B = rand(Float32, 2^10, 2^10, 2^10);
julia> AF_A = AFArray(A); AF_B = AFArray(B);
julia> @time batched_mul(A,B);
  3.167301 seconds (64 allocations: 4.000 GiB, 0.06% gc time)
julia> @time AF_A * AF_B;
  1.750176 seconds (1 allocation: 16 bytes, 2.09% gc time)

こう見るとCPUも早いですね... Float32 だとちょうど4GiBのメモリを食う巨大な計算です。

FFT

FFTも脳死で可能です。

julia> using FFTW
julia> A = rand(Float32, 2^25);
julia> @time fft(A);
  0.633331 seconds (29 allocations: 512.002 MiB, 2.18% gc time)
julia> AF_A = AFArray(A);
julia> @time fft(AF_A);
  0.000441 seconds (1 allocation: 16 bytes)

やっとGPUらしい高速化が見られましたね。

まとめ

ArrayFire+ OpenCL がApple Siliconに対応してくれているので、かなりの関数が既に実装されていて、CUDAに近い利便性があるのではないでしょうか。しかし、行列積周りはまだ最適ではないので数理最適化などのアプリケーションで速度を求める場合はMetal を勉強せざるを得ないのかなという印象です。

メモリ周りでApple Silicon は他のディスクリートGPUよりかなり有利なので、小規模な並列化ではApple Silicon は間違いなく強い選択肢の一つになってくると思います。また、もっと大規模なものでもヘテロジニアスコンピューティングは流行る可能性が高いと個人的には思っています。こういうところの開発はJulia の設計思想がとても有益なのでぜひ試してみていただきたいです。