4
5

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Day 15

Julia Advent Calendar 2025

@cometscome_phys(Yuki Nagai)

色々なGPUが使いたい in Julia

4
Last updated at Posted at 2025-12-14

JuliaでGPU計算、やってみたいですよね。そして、どうせGPUを使うなら、色々なGPUを使いたいですよね。
さらに、どうせGPU用のコードを書くなら、一回コードを書いたらそれがどのアーキテクチャでも動いてくれるといいですよね。NVIDIA製GPUでも、AMD製GPUでも、ついでに普通の並列も動いてくれるといいですよね。
そんなことできるの？とお思いの方に、JACC.jlというパッケージを紹介します。
このパッケージを使うと、自分の書いたコードが、NVIDIA製GPUでも、AMD製GPUでも、Intel製GPUでも、そして、CPUのマルチコアでも、並列に動いてしまいます。

インストール

インストールは、おすすめとしては、これをaddする前に、]キーを押してパッケージモードにしてから、

activate .
add JACC

です。これで、今コードを書こうとしているディレクトリ専用の仮想環境が用意され、そこにJACCが入ります。activate .は便利です。Pythonのvenvみたいなものですね。JACCは2025年12月7日現在0.6となっています。

次に、どのハードウェアを使うか、バックエンドを指定するために、

using JACC
JACC.set_backend("cuda")

のようにしてください。cudaだとNVIDIA製、amdgpuだとAMD製、threadsだとマルチコアです。
そのあとREPLを一度終了してください。これによって、LocalPreferences.tomlというファイルがディレクトリにできていると思います。中身は

[JACC]
backends = ["threads", "cuda"]
default_backend = "cuda"

    [JACC.placement]
    cuda = "none"

こんな感じです。

対象とする計算

次に、GPU計算のコードを書いてみます。
JACCを使う場合、どのハードウェアを使うかという関連の初期化をするために、最初に

julia
using JACC
JACC.@init_backend

とします。

さて、今回は、

C_{ij}(n) = \sum_{k=1}^{NC} A_{ik}(n)B_{kj}(n)

という計算をやってみることにします。ここで、インデックスnは実空間のインデックスで、つまり、n番目の格子点の上に$NC \times NC$の行列が定義されており、その行列の積を計算します。
この計算をGPUでやるにあたって、インデックスnの部分を並列化します。それぞれのnで計算は別々にできますので。ということで、以下のような関数を定義します。

@inline function kernel_mul!(n, C, A, B, NC)
    for jc = 1:NC
        for ic = 1:NC
            C[ic, jc, n] = 0.0
            for kc = 1:NC
                C[ic, jc, n] += A[ic, kc, n] * B[kc, jc, n]
            end
        end
    end
    return
end

最初の引数は$n$で、n番目の格子を意味しています。そして、C,A,Bは

    NC = 5
    N = 1000000
    C = zeros(Float64, NC, NC, N)
    A = rand(Float64, NC, NC, N)
    B = rand(Float64, NC, NC, N)

のように定義します。ここで1000000個の格子点の上に5x5行列が定義されています。
普通にCPUで計算するのであれば、

function multCPU!(C, A, B, NC, N)
    for n = 1:N
        kernel_mul!(n, C, A, B, NC)
    end
end

ですね。

JACCの書き方

次に、JACCを使って並列計算をしてみます。JACCでは、ハードウェアが何であるかを問わない書き方ができて、例えば、

    NC = 5
    N = 1000000
    C = zeros(Float64, NC, NC, N)
    A = rand(Float64, NC, NC, N)
    B = rand(Float64, NC, NC, N)

    Cj = JACC.array(C)
    Aj = JACC.array(A)
    Bj = JACC.array(B)

とします。JACC.array(C)とした瞬間、もしGPUを使っている場合には、配列CがGPUメモリーに転送されることになります。マルチCPUの場合はそのままです。

並列計算は、

function multGPU!(C, A, B, NC, N)
    JACC.parallel_for(
        N, kernel_mul!, C, A, B, NC
    )
end

です。最初の引数は、ループの長さを、第二引数は関数を、第三以降はその関数の引数を並べます。これで、kernel_mul!の第一引数を並列化することになります。

以上です。このコードで、どんなGPUでも並列化できてしまいます。

全体のコードは

using JACC
JACC.@init_backend

@inline function kernel_mul!(n, C, A, B, NC)
    for jc = 1:NC
        for ic = 1:NC
            C[ic, jc, n] = 0.0
            for kc = 1:NC
                C[ic, jc, n] += A[ic, kc, n] * B[kc, jc, n]
            end
        end
    end
    return
end

function multGPU!(C, A, B, NC, N)
    JACC.parallel_for(
        N, kernel_mul!, C, A, B, NC
    )
end

function multCPU!(C, A, B, NC, N)
    for n = 1:N
        kernel_mul!(n, C, A, B, NC)
    end
end

function main()
    # Your main application logic goes here
    NC = 5
    N = 1000000
    C = zeros(Float64, NC, NC, N)
    A = rand(Float64, NC, NC, N)
    B = rand(Float64, NC, NC, N)

    Cj = JACC.array(C)
    Aj = JACC.array(A)
    Bj = JACC.array(B)

    multGPU!(Cj, Aj, Bj, NC, N)
    display(Cj[:, :, 1])
    multCPU!(C, A, B, NC, N)
    display(C[:, :, 1])

    C = zeros(Float64, NC, NC, N)
    A = rand(Float64, NC, NC, N)
    B = rand(Float64, NC, NC, N)

    Cj = JACC.array(C)
    Aj = JACC.array(A)
    Bj = JACC.array(B)

    @time multGPU!(Cj, Aj, Bj, NC, N)
    display(Cj[:, :, 1])
    @time multCPU!(C, A, B, NC, N)
    display(C[:, :, 1])

end

main()

です。CPUとJACCの比較をしています。

ベンチマーク

それでは、計算がちゃんと並列化されているかみてみましょう。

スレッド並列

まず、マルチコアCPUを使ってみます。使うものはMacBookPro 14インチ(2023)、CPUはApple M2 Maxで、12コア（パフォーマンスコア8コア、効率化コア4コア）です。

スクリーンショット 2025-12-07 15.09.06.png

Juliaでスレッド並列をするには、-tをつけます。先ほどのディレクトリの環境を使うには、--project=.とします。
つまり、

julia -t 2 --project=. main.jl

で2スレッド並列ですね。
実行結果は

  0.129510 seconds (12 allocations: 1.141 KiB)
5×5 Matrix{Float64}:
 1.05593   0.953626  0.525468  1.03582   1.05768
 1.46264   2.52129   1.65358   1.38293   2.02142
 0.855354  1.10988   0.747243  0.669073  1.19669
 0.31872   0.470203  0.364287  0.323551  0.381042
 0.668523  1.09341   0.501557  0.68106   0.855375
  0.213086 seconds
5×5 Matrix{Float64}:
 1.05593   0.953626  0.525468  1.03582   1.05768
 1.46264   2.52129   1.65358   1.38293   2.02142
 0.855354  1.10988   0.747243  0.669073  1.19669
 0.31872   0.470203  0.364287  0.323551  0.381042
 0.668523  1.09341   0.501557  0.68106   0.855375

で大体2倍速になっています。4スレッドだと

  0.062821 seconds (22 allocations: 2.219 KiB)
5×5 Matrix{Float64}:
 1.25792   1.32799   0.940223  0.781699  1.55357
 0.878267  1.17095   0.992576  0.669238  1.12706
 1.20005   1.18871   0.607388  0.556967  1.44924
 0.562391  1.14431   0.880372  0.638255  0.804898
 0.292237  0.692637  0.408039  0.250453  0.316187
  0.205496 seconds
5×5 Matrix{Float64}:
 1.25792   1.32799   0.940223  0.781699  1.55357
 0.878267  1.17095   0.992576  0.669238  1.12706
 1.20005   1.18871   0.607388  0.556967  1.44924
 0.562391  1.14431   0.880372  0.638255  0.804898
 0.292237  0.692637  0.408039  0.250453  0.316187

8スレッドだと、

  0.027700 seconds (42 allocations: 4.375 KiB)
5×5 Matrix{Float64}:
1.22701  0.96939  0.506475  0.489596  1.16823
2.23987  1.83774  0.572637  1.05251   1.89627
1.67783  1.38032  0.404913  1.03638   1.52376
1.55576  1.05758  0.40458   0.998169  1.32818
1.95967  1.66258  0.589071  0.606719  1.62389
 0.186760 seconds
5×5 Matrix{Float64}:
1.22701  0.96939  0.506475  0.489596  1.16823
2.23987  1.83774  0.572637  1.05251   1.89627
1.67783  1.38032  0.404913  1.03638   1.52376
1.55576  1.05758  0.40458   0.998169  1.32818
1.95967  1.66258  0.589071  0.606719  1.62389

で、8倍速にはなっていませんが結構速くなっています。なお、スレッドを使うと、

[ Info: Threads backend loaded with 8 threads

のようなメッセージが出ますので、８スレッドで動いていることがわかります。

NVIDIA GeForce RTX 4090

次に、RTX4090を使ってみましょう。
スクリーンショット 2025-12-07 14.08.21.png
結果は、

  0.006675 seconds (74 allocations: 3.891 KiB)
5×5 CUDA.CuArray{Float64, 2, CUDA.DeviceMemory}:
 1.17201   1.59068   1.72769   1.99373  1.59829
 0.998139  1.32721   1.3987    1.49103  1.54037
 1.8172    1.79746   1.7754    2.86208  1.72581
 0.972047  1.48182   1.53006   1.77494  1.38059
 0.609081  0.994494  0.867954  1.39835  0.772122
  0.291615 seconds
5×5 Matrix{Float64}:
 1.17201   1.59068   1.72769   1.99373  1.59829
 0.998139  1.32721   1.3987    1.49103  1.54037
 1.8172    1.79746   1.7754    2.86208  1.72581
 0.972047  1.48182   1.53006   1.77494  1.38059
 0.609081  0.994494  0.867954  1.39835  0.772122

です。速いですね。倍精度演算なので単精度の1/64しか速度が出ないはずですが、今回みたいな小さな行列を何回も計算するやつは、多分GPUメモリ律速になっていると思われます。
また、CUDAが使われている場合には、

[ Info: CUDA backend loaded

と出ます。これが出ていないと、GPUが使えていませんので注意してください。

AMD Radeon RX 7800 XT

次は、AMD製のGPUです。コードは全く同じで構いません。
スクリーンショット 2025-12-07 14.11.30.png
実行結果は、

  0.029817 seconds (334 allocations: 9.016 KiB)
5×5 AMDGPU.ROCArray{Float64, 2, AMDGPU.Runtime.Mem.HIPBuffer}:
 0.37837   0.446957  0.590148  0.72041  0.747445
 0.552329  0.615178  0.729361  1.02719  0.71334
 1.15405   0.673195  0.86338   1.78428  1.31944
 1.02976   0.759977  1.3385    2.15536  0.703746
 1.16309   0.899828  1.27847   2.23525  0.78781
  0.153014 seconds
5×5 Matrix{Float64}:
 0.37837   0.446957  0.590148  0.72041  0.747445
 0.552329  0.615178  0.729361  1.02719  0.71334
 1.15405   0.673195  0.86338   1.78428  1.31944
 1.02976   0.759977  1.3385    2.15536  0.703746
 1.16309   0.899828  1.27847   2.23525  0.78781

です。RTXと比べるとメモリ帯域が確か半分だったので、メモリ律速を裏付けているような感じです。
AMDの場合は、実行すると、

[ Info: AMDGPU backend loaded

と出ます。

NVIDIA H100 NVL

最後に、NVIDIA H100 NVLが手元で使えたので、こちらでも試してみます。

スクリーンショット 2025-12-07 14.08.01.png

実行結果は、

  0.007441 seconds (323 allocations: 7.547 KiB)
5×5 CUDA.CuArray{Float64, 2, CUDA.DeviceMemory}:
 0.615174  1.01216  1.25659  1.30054  0.872642
 1.47844   1.6595   2.11988  2.18352  2.0313
 0.963653  1.17193  1.48219  1.6469   1.15857
 1.4952    1.77438  1.98736  1.96495  1.3483
 0.992886  1.19461  1.4507   1.06382  1.55244
  0.233601 seconds
5×5 Matrix{Float64}:
 0.615174  1.01216  1.25659  1.30054  0.872642
 1.47844   1.6595   2.11988  2.18352  2.0313
 0.963653  1.17193  1.48219  1.6469   1.15857
 1.4952    1.77438  1.98736  1.96495  1.3483
 0.992886  1.19461  1.4507   1.06382  1.55244

です。RTXとあまり変わらない結果というのは、やはりGPUメモリ律速なのかなと思います。

4
5
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
4
5

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?