JuliaでCUDA.jlつかって行列積と固有値分解してみた

環境

GeForce RTX3090，Xeon Gold 6338 (2.0GHz)搭載のRocky LinuxマシンでCUDA.jl使って行列計算してみました。

julia> versioninfo()
Julia Version 1.7.2
Commit bf53498635 (2022-02-06 15:21 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, icelake-server)

CUDAをインストールした以外は特に変わったことはしてません。

行列積

CPUとGPU計算を比較。倍精度と単精度を両方比較。

using CUDA
using BenchmarkTools
using LinearAlgebra

N = 1000
dim = N


println("Float64")
A = randn(Float64, dim, dim)
B = randn(Float64, dim, N)
C = zeros(Float64, dim, N)

AC = CUDA.randn(Float64, dim, dim)
BC = CUDA.randn(Float64, dim, N)
CC = CUDA.zeros(Float64, dim, N)

@btime mul!(C, A, B)
@btime CUDA.@sync mul!(CC, AC, BC)


println("Float32")
A = randn(Float32, dim, dim)
B = randn(Float32, dim, N)
C = zeros(Float32, dim, N)

AC = CUDA.randn(Float32, dim, dim)
BC = CUDA.randn(Float32, dim, N)
CC = CUDA.zeros(Float32, dim, N)

@btime mul!(C, A, B)
@btime CUDA.@sync mul!(CC, AC, BC)

計算時間

Float64
  28.881 ms (0 allocations: 0 bytes)
  3.981 ms (71 allocations: 3.33 KiB)
Float32
  15.027 ms (0 allocations: 0 bytes)
  121.564 μs (27 allocations: 464 bytes)

実用的には，GPU上のデータをCPU側にもってくる時間がGPUの場合必要になると思われますが，それでもGPUのほうがご利益多そう。倍精度でも10倍近く早いです。単精度だと，100倍ちかい。

固有値分解

固有値と固有ベクトルを計算します。注意としては，CUSOLVER.syevd!は対称行列しか使えない，引数として与えた配列は固有ベクトルを与えたものに置き換わります。使用時にご注意下さい。

using CUDA
using LinearAlgebra
using BenchmarkTools

N = 2000

ac = rand(Float64, N, N)
a = CuArray(ac)

println("Float64")
@btime eigen!(Symmetric(ac))
@btime CUSOLVER.syevd!('V','U',a)


println("Float32")

ac = rand(Float32, N, N)
a = CuArray(ac)

@btime eigen!(Symmetric(ac))
@btime CUSOLVER.syevd!('V','U',a)

計算時間

Float64
  1.926 s (14 allocations: 61.75 MiB)
  141.056 ms (70 allocations: 4.03 KiB)
Float32
  1.705 s (14 allocations: 30.97 MiB)
  37.263 ms (23 allocations: 1.11 KiB)

倍精度で10倍，単精度で46倍早くなりました。

（補足）
行列サイズが小さいとCPUのほうが計算が早くなります。N=4はCPUが早かったです。N=400はGPUのほうが早かったです。