JuliaとCUDAのシンプルな使い方

Julia言語でCUDAを使う方法を忘備録として残しておきます。

• Julia 1.11

で試しています。単純な計算なので、CUDAのバージョンはよほど古くない限り大丈夫でしょう。CUDA.jlのドキュメントを見ると色々書き方が書いてあります。

ある二次元格子上の座標点$(x,y)$において、$N \times N$行列が定義されているとします。そして、

C(x,y) = A(x,y) B(x,y) 

が計算したいとします。座標点ごとに計算は独立ですので、GPUでまとめて計算できるはずです。

コード例

using CUDA

function cudakernel_mul!(C,A,B,N)
    b = Int64(CUDA.threadIdx().x)
    r = Int64(CUDA.blockIdx().x)
    kernel_mul!(C,A,B,N,b,r)
end

function kernel_mul!(C,A,B,N,b,r)
    for j=1:N
        for i=1:N
            C[i,j,b,r] = 0
            for k=1:N
                C[i,j,b,r] += A[i,k,b,r]*B[k,j,b,r]
            end
        end
    end
end

function testcudamul!(C,A,B,N,blocksize,rsize) 
    CUDA.@sync begin
        CUDA.@cuda threads = blocksize blocks = rsize cudakernel_mul!(C,A,B,N)
    end
end

function testmul!(C,A,B,N,blocksize,rsize) 
    for r=1:rsize
        for b=1:blocksize
            kernel_mul!(C,A,B,N,b,r)
        end
    end
end

function main()
    rsize = 64
    blocksize = 256
    N = 3
    A = rand(N,N,blocksize,rsize)
    B = rand(N,N,blocksize,rsize)
    C = zeros(N,N,blocksize,rsize)

    Agpu = CuArray(A)
    Bgpu = CuArray(B)
    Cgpu = CuArray(C)

    testmul!(C,A,B,N,blocksize,rsize) 
    display(C[:,:,4,5])

    testcudamul!(Cgpu,Agpu,Bgpu,N,blocksize,rsize) 
    display(Cgpu[:,:,4,5])

end

main()

です。
これだけです。つまり、各ループの中で独立に計算できるものは、GPUで簡単に並列で計算できます。