Adding GPU support, exploring options

[raar1]

We have discussed adding GPU support to PartitionedArrays. There are a number of possibilities for doing this, with the most promising being (naturally) trying to wrap existing GPU solvers as much as possible. Nevertheless, there is always the question of vendor specificity/agnosticism and maintainability.

[raar1]

I add below an adaptation of the thoughts of @Cmurilochem on this:

Julia has an entire ecosystem devoted to gpus: https://github.com/juliagpu.
Among the several wrappers, CUDA.jl (https://github.com/JuliaGPU/CUDA.jl) appears to be the most mature so far. Depending of what we want and how our software is designed, the CUDA.jl allow us to levarage GPUs using wither a high-level or low-level. The high-level implementation concerns the conversion of any allocated array (in the host/cpu) to an special CuArray type that os automatically transferred to the GPU (device) by CUDA.jl. Of course, this is not that simple and we need to also write specific kernels and can go low-level as CUDA C itself. An example implementation is here: https://github.com/QuantumBFS/Yao.jl. Here the guy implements CUDA.jl as an extension module.
Apart from CUDA, there is also a very popular package that allows the implementation of abstract kernels, agnostic to specific vendors. This one looks very nice but the docs are very poor. It is called KernelAbstractions.jl (https://github.com/JuliaGPU/KernelAbstractions.jl). An example implementation can be found here: https://github.com/SciML/DiffEqGPU.jl. Now, this package is completely separated from it cpu-counterpart and the extension modules are those wrappers for the specific vendors.
I've tried to make a simple test for PartiotionedSolvers but as expected it looks like this is much more than just converting arrays to CuArrays. The problem is much more complicated than that and will require us writing the specific kernels.

[raar1]

Note that KernelAbstractions is also what we have been using in our Efficient Julia course at NLeSC: https://carpentries-incubator.github.io/efficient-computing-in-julia/170-gpu-programming.html

I certainly agree with Carlos about its documentation being incredibly sparse...

However it is nice to be able to use KernelAbstractions for any GPU specific code, and then users need only install the backend relevant to their architecture. e.g.

• using oneAPI for Intel
• using CUDA for NVidia
• using AMDGPU for AMD
• using Metal for Macs

[fverdugo]

Here there is a working example:

https://discourse.julialang.org/t/sparse-lu-factorization-on-gpu/122012/4?u=fverdugo

It is pretty much what we need to start.

[Cmurilochem]

Thanks @raar1 and @fverdugo for example. I was curious and tried to implement a minimal example using PartitionedSolvers. Here it is:

import PartitionedSolvers as PS
using CUDA

using BenchmarkTools
using Test

CUDA.allowscalar(false)

function mock_linear_solver(p)
    @assert ! PS.uses_mutable_types(p)
    Ainv = inv(PS.matrix(p))
    function update(Ainv,A)
        Ainv = inv(A)
    end
    function step(x,Ainv,b,phase=:start)
        x = Ainv*b
        phase = :stop
        x,Ainv,phase
    end
    uses_initial_guess = false
    PS.linear_solver(update,step,p,Ainv;uses_initial_guess)
end

N = 10^4
x = zeros(N)
A = rand(N,N)
b = rand(N)


lp = PS.linear_problem(x,A,b;uses_mutable_types=false)
ls = mock_linear_solver(lp)
ls = PS.solve(ls)
x_cpu = PS.solution(ls)
@test x_cpu ≈ A\b

@btime begin
    ls = mock_linear_solver($lp)
    ls = PS.solve(ls)
    PS.solution(ls)
end

x_d = zeros(N) |> cu
A_d = A |> cu
b_d = b |> cu

@test x_d isa CuArray
@test A_d isa CuArray
@test b_d isa CuArray

lp_d = PS.linear_problem(x_d,A_d,b_d;uses_mutable_types=false)
ls_d = mock_linear_solver(lp_d)
ls_d = PS.solve(ls_d)
x_gpu = PS.solution(ls_d)

@test Array(x_gpu) ≈ x_cpu

@btime begin
    CUDA.@sync begin
        ls_d = mock_linear_solver($lp_d)
        ls_d = PS.solve(ls_d)
        PS.solution(ls_d)
    end
end

I tested using the GPU (int3) iterative node in snellius and the GPU is much faster indeed by just using CuArrays.