はじめに

(2022/7/13 追記:1ヶ月余りでかなり開発が進んでおり、記載がかなり古くなっています。ご注意ください。)

Julia 言語はCUDA.jl をはじめとして、GPUを使用した行列計算が容易にできるプラットフォームとしても有用です。

Apple SiliconのGPUを使えるAPI としては、OpenCLとMetalが有力候補で、AppleはMetalの利用を推奨しています。つい先日(2022/5/24) のMetal.jl とGPUCompiler の更新でコード改変なしで Metal が利用できるようになった(っぽい)ので紹介したいと思います。

筆者の環境は

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

です。

Metal.jl のインストール

かなり開発途上といった感じで、不安定版のパッケージに依存しています。

必要条件

M1 チップのMac
Xcode Command Line Tool
Julia 1.8
LLVM v4.13.0
GPUCompiler v0.16.0

インストール

2022/7/13 追記: Pkgのmain環境に]add Metal のみでパッケージが追加できるようになりました。
いくつかのパッケージとの衝突が残っているようですので、プロジェクト環境を利用する下のやり方を現時点でもまだ推奨します。

Metal.jl のReadMe には
"Bleeding Edge" なパッケージに多く依存するので、 --project パラメータをつけて実行することを推奨する
とあります。~~この場合一部パッケージのパスに修正が必要 (2022/5/31) で面倒なので、~~ (2022/07/13 修正を確認)

Pkg のプロジェクト環境を利用するなどしましょう。
現状 Flux.jl などのパッケージと衝突する模様です。

$ git clone https://github.com/JuliaGPU/Metal.jl && cd Metal.jl
$ julia
pkg> activate .
pkg> build

テスト

まずは、公式のテストから。

pkg> test Metal
┌ Info: System information:
│ macOS 12.2.0, Darwin 21.3.0
│ 
│ Toolchain:
│ - Julia: 1.8.0-beta3
│ - LLVM: 13.0.1
│ 
│ 1 device:
└ - Apple M1 Max (64.000 KiB allocated)
                |          | ---------------- CPU ---------------- |
Test   (Worker) | Time (s) | GC (s) | GC % | Alloc (MB) | RSS (MB) |
array       (2) |     0.17 |   0.00 |  2.4 |      44.85 |   613.30 |
core        (2) |     0.50 |   0.00 |  0.9 |      52.63 |   629.55 |
examples    (2) |    46.60 |   0.00 |  0.0 |      24.64 |   632.50 |
execution   (2) |    13.95 |   0.43 |  3.1 |    2303.83 |   751.34 |
intrinsics  (2) |     2.76 |   0.13 |  4.6 |     746.84 |   773.67 |
Testing finished in 1 minute, 6 seconds, 113 milliseconds

Test Summary: | Pass  Total  Time
  Overall     |  120    120      
    SUCCESS

簡易的なFLOPS の評価も可能です。M1 Max 24 Core は大体8TFlopsらしい(https://appleinsider.com/articles/21/10/19/m1-pro-and-m1-max-gpu-performance-versus-nvidia-and-amd) のでそれっぽい

julia> using Metal
julia> include("examples/peakflops.jl")
TFlops: 7.21

Metal.jl を利用した計算

Array Abstraction

JuliaGPU ではお馴染みのGPUArray (CuArray etc..) を利用した計算が一部可能です。例えば

julia> a = MtlArray(rand(Float32,10))
10-element MtlArray{Float32, 1}:
 0.47226274
 0.18609709
 0.57571167
 0.98504615
 0.18164462
 0.18123078
 0.085390806
 0.025278926
 0.5564378
 0.29086506

julia> a .+ 1.0f0
10-element MtlArray{Float32, 1}:
 1.4722627
 1.1860971
 1.5757117
 1.9850461
 1.1816447
 1.1812308
 1.0853908
 1.0252789
 1.5564377
 1.2908651

しかしArray Abstraction はまだエラーが多いです。(2022/7/13) ~~ドット演算子もベクトルにしか対応していない模様.~~

(2022/07/13) 修正を確認. 現状行列積および、初等的な関数のドット演算子に対応を確認。しかし、.^2ができて .^4はできないなどバグまだ多いです。

julia> Ndims = 2^10;
julia> a = rand(Float32,Ndims,Ndims); b = rand(Float32,Ndims,Ndims);
julia> Mt_a = MtlArray(a); Mt_b = MtlArray(b);
julia> @time a *b;
0.015371 seconds (2 allocations: 4.000 MiB)
julia> @time Mt_a * Mt_b
  0.001089 seconds (237 allocations: 5.828 KiB)

そして、超早くなってます。最高！

Kernel Programming

(2022/07/13 加筆) Kernel から書く場合は以下のようになります。
特に、二次元配列に対する演算の記法に関してはドキュメントはない状況ですから、下の表記を参考にしてみてください。


julia> Ndims = 2^10;

julia> Ngrids = floor(Int,Ndims/32);

julia> a = rand(Float32,Ndims,Ndims); b = rand(Float32,Ndims,Ndims);

julia> Mt_a = MtlArray(a); Mt_b = MtlArray(b);

julia> Mt_c = similar(Mt_a);

julia> function matmul(c, a, b, Ndims)
                  i = thread_position_in_grid_2d().x
                  j = thread_position_in_grid_2d().y
                  for k = 1:Ndims
                      c[i,j] += a[i,k] * b[k,j]
                  end
                  return
              end
matmul (generic function with 1 method)

julia> @time Metal.@sync @metal threads=(32,32) grid = (Ngrids,Ngrids) matmul(Mt_c, Mt_a, Mt_b, Ndims);
  0.050729 seconds (36.79 k allocations: 1.769 MiB, 11.49% compilation time)

julia> @time a*b;
  0.015356 seconds (2 allocations: 4.000 MiB)

julia> Array(Mt_c) ≈ a * b
true

Thread Index の表記はCUDA とは大きく異なっている(らしい。)
~~そして、あんまり早くない~~(Kernel を適当に書いているので最適化はされていない)

Unified Memory

ここからが本題かもしれません。
M1 GPU はCPUとGPU がメモリを共有しているためにCPU-GPU間でシームレスに計算ができることが特徴です。

unsafe_wrapにより同じプールにあるGPU Array をCPU Array として扱うことができます。

julia> A = [rand(Float32)]
1-element Vector{Float32}:
 0.27993953
julia> mt_A = MtlArray(A)
1-element MtlArray{Float32, 1}:
 0.27993953
julia> A_cpu = unsafe_wrap(Array{Float32}, mt_A, 1)
1-element Vector{Float32}:
 0.27993953
julia> mt_A .+= 1.0f0
1-element MtlArray{Float32, 1}:
 1.2799395
julia> A_cpu
1-element Vector{Float32}:
 1.2799395

このように、GPUArray であるmt_A に加算を行うとA_cpu に反映されます。さらに、

julia> A_cpu .= 1.0
1-element Vector{Float32}:
 1.0
julia> mt_A
1-element MtlArray{Float32, 1}:
 1.0

A_cpu を編集すると、mt_A にも反映されます。これはつまりGPUが不得意な配列の特定要素へのアクセスなどをCPU側で行ってネックにならないようにできることを意味します。

まとめ

Julia + M1 Mac でGPU が使えるようになりました。開発者のみなさまには頭が上がらないですね....
今後Appleによって最適化された行列演算ライブラリなどが実装されてくるとかなり盛り上がるのではないでしょうか。

Julia で Apple Silicon (M1) のGPU 計算 (Metal.jl)