JuliaのGPU並列化の簡素な手法

JuliaでGPUを用いた計算を並列化したい

重い計算をしなくてはならなくなりまして、どうにかGPUを用いたスレッド並列化がしたい！！！

方法

以下、実際のサンプルコードになります。
多分ですが、CUDAのカーネルまでいじればもっと能率的なコードが書けるんだとは思いますが、現在の僕の能力では、単純にスレッド並列化させるだけで精いっぱいでした。

パッケージ

次のパッケージを用いています。

using MKL
using CUDA
using LinearAlgebra

期待値の計算（サンプル）

function calc_expval(kx,ky)という関数の返り値(Float64)をすべて足した結果を得たい、という計算とします。

function calc_expval(kx,ky)
    H = Hamiltonian(kx,ky)
    X = Operator(kx,ky)

    # to GPU
    X_d = CuMatrix{ComplexF64}(X)
    H_d = CuMatrix{ComplexF64}(H)
    E_d, Ψ_d = CUSOLVER.heevd!('V', 'U', H_d)

    # Multiplication on GPU
    A_d = CuMatrix{ComplexF64}(undef,size(X_d,1),size(Ψ_d,2))
    B_d = CuMatrix{ComplexF64}(undef,size(Ψ_d,2),size(Ψ_d,2))
    CUBLAS.gemm!('N','N',
        1.0,X_d,Ψ_d,0.0,A_d
    ) # <=> A_d = X_d * Ψ_d
    CUBLAS.gemm!('C','N',
        1.0,Ψ_d,A_d,0.0,B_d
    ) # <=> B_d = (Ψ_d)' * (X_d * Ψ_d)

    # to CPU
    expval = Matrix{ComplexF64}(B_d)
    return tr(expval)
end

簡単な解説

この計算は、ハミルトニアンHamiltonian(kx,ky)から固有状態を計算して、観測量Operator(kx,ky)の期待値を計算することを目的としています。
CuMatrixを駆使して、GPUに配列（行列）を渡していき、可能な限りGPU側で計算を行っていきます。
CUDAに入っているCUSOLVERを用いて

E_d, Ψ_d = CUSOLVER.heevd!('V', 'U', H_d)

を計算することで、ハミルトニアンの固有状態Ψ_dが得られます。
あとは観測量を固有状態で挟めばよいわけですが、今回は次のような式で導きました。
$$
\braket{X} = \mathrm{Tr}\left[\Psi^{\dagger}X \Psi\right]
$$
以上の計算のうち、行列積はCUBLASを用いて次のように実装しました。

A_d = CuMatrix{ComplexF64}(undef,size(X_d,1),size(Ψ_d,2))
B_d = CuMatrix{ComplexF64}(undef,size(Ψ_d,2),size(Ψ_d,2))
CUBLAS.gemm!('N','N',
        1.0,X_d,Ψ_d,0.0,A_d
    ) # <=> A_d = X_d * Ψ_d
CUBLAS.gemm!('C','N',
        1.0,Ψ_d,A_d,0.0,B_d
    ) # <=> B_d = (Ψ_d)' * A_d

一つ目で$X$と$\Psi$の掛け算を評価したのちに、次に$\Psi^{\dagger}$と$(X\Psi)$の掛け算を行っています。本来は可能であれば、一度に(Ψ_d)' * X_d * Ψ_dと計算したいのですが、このように書くとなぜかGPUのガベージコレクションがうまく働いてくれず、あっという間にOut of memoryになるため、このように二回に分けて、用いるメモリを指定しています。

（ここは、公式のissueでも同様のことが書かれていたので、今後直る可能性があります。）

そして最後に、GPUのデータをCPUに戻してきて、returnしています。

# to CPU
expval = Matrix{ComplexF64}(B_d)
return tr(expval)

並列化の実装

続いて、この計算を波数kx,kyについてスレッド並列化させます。

function main_gpu_parallel()
  №_th = nthreads()

  kx_list = range(-π,π,length=kmesh+1)[1:end-1]
  ky_list = range(-π,π,length=kmesh+1)[1:end-1]
  k_pair = Iterators.product(kx_list,ky_list)

  expval_th = zeros(ComplexF64,№_th)

  @threads for (kx,ky) ∈ collect(k_pair)
    №_id = threadid()
    device!(№_id % Int(length(devices())))

    expval_th[№_id] += calc_expval(kx,ky)

    GC.gc()
  end

  expval = sum(expval_th,dims=1)[1]
  
  return expval
end

簡単な解説

今回は、波数の配列k*_listをIterators.productですべての組について用意しておくことで、k_pairに関するループを回すようにしています。というのも、例えばkmesh=32程度だった場合に、それらを64並列すると無駄になってしまうからです。組を取っておけば、32^2要素を64で並列化できるため、無駄なく並列化できると考えました。

また、№_th = nthreads()で並列化のスレッド数を取得し、その個数分のデータ保存場所としてexpval_th = zeros(ComplexF64,№_th)を定義しています。こちらもexpval_k = zeros(ComplexF64,kmesh,kmesh)のようにしてそれぞれの波数に対する保存を行ってもよいのですが、今回ほしいのは和だけですので、配列の無駄のないように最低限の要素にしています。

次は複数GPUがある場合の処方箋です。スレッド並列のループ内で

№_id = threadid()
device!(№_id % Int(length(devices())))

という処理をしております。こちらは、device!(№_id % Int(length(devices())))によって、用いるGPUの番号を指定しています。そうすることで、複数のGPUがある計算機でも、すべてのGPUを用いて計算してくれるようになるはずです。

最後に

今回はシンプルにGPUの計算をスレッド並列化させるという内容でした。これだけでもかなり使い物になったので、ぜひ参考にして下さい。

また、本来はもっといじれば性能よく並列化できそうな気もしているので、引き続き勉強します！