Pure Juliaの行列演算ライブラリGaius.jlでIntel MKLを倒したい
Qiitaに書いた記事をこちらにも載せておきます。内容は全く同じです。
LoopVectorization.jlを使うとPure JuliaでIntel MKLに匹敵するスピードの行列演算ができるとのことで コミュニティが盛り上がっているようです。
さらにLoopVectorization.jlを使ったBLAS-likeライブラリである Gaius.jlも作られているようなので、早速試してみました。
LoopVectorization.jlとは
Juliaのfor
ループをベクトル化してくれる@avx
というマクロを提供してくれます。公式のExampleに倣って
function my_gemm!(Y, A, B)
@avx for m in 1:size(A,1), n in 1:size(B,2)
Ymn = zero(eltype(Y))
for k in 1:size(A,2)
Ymn += A[m,k] * B[k,n]
end
Y[m,n] = Ymn
end
end
のように書けば、ベクトル化された行列行列積関数を実装することができます。
Gaius.jlとは
LoopVectorization.jlによるベクトル化された行列行列積関数と分割統治を組み合わせて作られた、BLAS-likeライブラリです。インストールは
]add https://github.com/MasonProtter/Gaius.jl
とすればOK。juliaのバージョンが1.3以上であることが要求されるので注意しましょう。
使い方はこんな感じ。
using Gaius
d = 1000
A, B = randn(d,d), randn(d,d)
C = similar(A)
Gaius.mul!(C, A, B) # Gaius.jlを使って計算
using LinearAlgebra
D = similar(C)
mul!(D, A, B) # Julia組み込みの行列行列積
C ≈ D # -> true
ちゃんと環境変数JULIA_NUM_THREADS
を使ってスレッド数を指定しておかないと性能が出ないので注意。スレッド数を変更した後はコンパイルし直す必要がある気がします。
性能比較
Gaius
とLinearAlgebra
のmul!
の性能を比較してみましょう。
今回テストを行う環境は以下の通りです。
Julia Version 1.3.1
Platform Info:
OS: Linux (x86_64-pc-linux-gnu)
"CentOS release 6.5 (Final)"
CPU: Intel(R) Xeon(R) CPU E5-2670 0 @ 2.60GHz:
speed user nice sys idle irq
#1-16 1200 MHz 9233372591 s 8111 s 147348681 s 8423904558 s 1 s
Memory: 62.743385314941406 GB (60272.515625 MB free)
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-6.0.1 (ORCJIT, sandybridge)
Environment:
JULIA_NUM_THREADS = 16
組み込みのmul!
はMKLのを使います。
julia> BLAS.vendor()
:mkl
実験結果
サイズNの正方行列をランダム生成して、BenchmarkTools
を使って行列行列積にかかる時間を測定します。
使用したコードは以下の通り。
using BenchmarkTools, Gaius, LinearAlgebra
results_mkl, results_gaius = BenchmarkTools.Trial[], BenchmarkTools.Trial[]
for n in 100: 100:5000
A, B, Y = randn(n,n), randn(n,n), zeros(n,n)
t_mkl = @benchmark mul!($Y, $A, $B)
t_gaius = @benchmark Gaius.mul!($Y, $A, $B)
push!(results_mkl, t_mkl)
push!(results_gaius, t_gaius)
end
結果をプロットしてみます。
頑張っていますが、流石にMKLを倒すのは厳しいですね。MKLのほうが2.5倍ほど速いです。思ったよりもボロ負けでした。他の環境でどうなるかも気になるところです。
Float64
以外での比較やOpenBLASとの比較もしたかったですが、面倒だったので省略しています。(誰かやってください。)
まとめ
LoopVectorization.jlとGaius.jlの紹介とベンチマークをしました。今後に期待しましょう。他の環境で試してみた報告や、ここを変えると速くなるよ報告等、お待ちしています。