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

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

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

## パッケージ
次のパッケージを用いています。
```julia
using MKL
using CUDA
using LinearAlgebra
```

## 期待値の計算（サンプル）
`function calc_expval(kx,ky)`という関数の返り値(`Float64`)をすべて足した結果を得たい、という計算とします。
```julia
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`を用いて
```julia
E_d, Ψ_d = CUSOLVER.heevd!('V', 'U', H_d)
```
を計算することで、ハミルトニアンの固有状態`Ψ_d`が得られます。
あとは観測量を固有状態で挟めばよいわけですが、今回は次のような式で導きました。
$$
    \braket{X} = \mathrm{Tr}\left[\Psi^{\dagger}X \Psi\right]
$$
以上の計算のうち、行列積は`CUBLAS`を用いて次のように実装しました。
```julia
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](https://discourse.julialang.org/t/multi-threaded-calls-to-cuda-matrix-multiplication/102753)でも同様のことが書かれていたので、今後直る可能性があります。）

そして最後に、GPUのデータをCPUに戻してきて、`return`しています。
```julia
# to CPU
expval = Matrix{ComplexF64}(B_d)
return tr(expval)
```

## 並列化の実装
続いて、この計算を波数`kx,ky`についてスレッド並列化させます。
```julia
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がある場合の処方箋です。スレッド並列のループ内で
```julia
№_id = threadid()
device!(№_id % Int(length(devices())))
```
という処理をしております。こちらは、`device!(№_id % Int(length(devices())))`によって、用いるGPUの番号を指定しています。そうすることで、複数のGPUがある計算機でも、すべてのGPUを用いて計算してくれるようになるはずです。


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

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