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の紹介とベンチマークをしました。今後に期待しましょう。他の環境で試してみた報告や、ここを変えると速くなるよ報告等、お待ちしています。