Tentative CUDA/AMGX support

Project.toml

@@ -1,7 +1,7 @@
 name = "JutulDarcy"
 uuid = "82210473-ab04-4dce-b31b-11573c4f8e0a"
 authors = ["Olav Møyner <olav.moyner@gmail.com>"]
-version = "0.2.35"
+version = "0.2.36"
 
 [deps]
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
@@ -34,20 +34,26 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [weakdeps]
+AMGX = "c963dde9-0319-47f5-bf0c-b07d3c80ffa6"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 
 [extensions]
+JutulDarcyAMGXExt = "AMGX"
+JutulDarcyCUDAExt = "CUDA"
 JutulDarcyGLMakieExt = "GLMakie"
 JutulDarcyMakieExt = "Makie"
 JutulDarcyPartitionedArraysExt = ["PartitionedArrays", "MPI", "HYPRE"]
 
 [compat]
 AlgebraicMultigrid = "0.5.1, 0.6.0"
 Artifacts = "1"
+AMGX = "0.2"
+CUDA = "5"
 DataStructures = "0.18.13"
 Dates = "1"
 DelimitedFiles = "1.6"
@@ -56,7 +62,7 @@ ForwardDiff = "0.10.30"
 GLMakie = "0.10.13"
 GeoEnergyIO = "1.1.12"
 HYPRE = "1.6.0, 1.7"
-Jutul = "0.2.40"
+Jutul = "0.2.41"
 Krylov = "0.9.1"
 LazyArtifacts = "1"
 LinearAlgebra = "1"
@@ -78,6 +84,8 @@ Tullio = "0.3.4"
 julia = "1.7"
 
 [extras]
+AMGX = "c963dde9-0319-47f5-bf0c-b07d3c80ffa6"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
@@ -87,4 +95,4 @@ TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [targets]
-test = ["Test", "TestItemRunner", "TestItems", "HYPRE", "MPI", "PartitionedArrays"]
+test = ["Test", "CUDA", "TestItemRunner", "TestItems", "HYPRE", "MPI", "PartitionedArrays"]

ext/JutulDarcyAMGXExt/JutulDarcyAMGXExt.jl

@@ -0,0 +1,15 @@
+module JutulDarcyAMGXExt
+    using JutulDarcy, Jutul, AMGX, AMGX.CUDA, LinearAlgebra, SparseArrays
+    import JutulDarcy: AMGXPreconditioner
+    import Jutul: @tic
+
+    timeit_debug_enabled() = Jutul.timeit_debug_enabled()
+
+    include("cpr.jl")
+    include("precond.jl")
+
+    function __init__()
+        # TODO: Figure out a way to not always do this?
+        AMGX.initialize()
+    end
+end

ext/JutulDarcyAMGXExt/cpr.jl

@@ -0,0 +1,106 @@
+function JutulDarcy.update_amgx_pressure_system!(amgx::AMGXPreconditioner, A_p::Jutul.StaticCSR.StaticSparsityMatrixCSR, Tv)
+    data = amgx.data
+    update_pressure_system!(amgx, A_p, Tv)
+    return amgx
+end
+
+mutable struct AMGXStorage{C, R, V, M, S}
+    config::C
+    resources::R
+    r::V
+    x::V
+    matrix::M
+    solver::S
+    function AMGXStorage(config, amgx_mode = AMGX.dDDI)
+        resources = AMGX.Resources(config)
+        r = AMGX.AMGXVector(resources, amgx_mode)
+        x = AMGX.AMGXVector(resources, amgx_mode)
+        matrix = AMGX.AMGXMatrix(resources, amgx_mode)
+        solver = AMGX.Solver(resources, amgx_mode, config)
+        function finalize_storage!(amgx_s::AMGXStorage)
+            @async println("Finalizing AMGXStorage")
+            close(amgx_s.solver)
+            close(amgx_s.r)
+            close(amgx_s.x)
+            close(amgx_s.matrix)
+            close(amgx_s.resources)
+            close(amgx_s.config)
+        end
+        s = new{
+            typeof(config),
+            typeof(resources),
+            typeof(r),
+            typeof(matrix),
+            typeof(solver)
+        }(config, resources, r, x, matrix, solver)
+        return finalizer(finalize_storage!, s)
+    end
+end
+
+function JutulDarcy.gpu_amgx_solve!(amgx::AMGXPreconditioner, r_p)
+    Jutul.apply!(r_p, amgx, r_p)
+end
+
+function update_pressure_system!(amgx::AMGXPreconditioner, A::Jutul.StaticCSR.StaticSparsityMatrixCSR, Tv)
+    if !haskey(amgx.data, :storage)
+        if Tv == Float64
+            amgx_mode = AMGX.dDDI
+        else
+            amgx_mode = AMGX.dFFI
+        end
+        n, m = size(A)
+        @assert n == m
+        config = AMGX.Config(amgx.settings)
+        s = AMGXStorage(config, amgx_mode)
+        # RHS and solution vectors to right size just in case
+        AMGX.set_zero!(s.x, n)
+        AMGX.set_zero!(s.r, n)
+
+        row_ptr = Cint.(A.At.colptr .- 1)
+        colval = Cint.(A.At.rowval .- 1)
+        nzval = A.At.nzval
+        # TODO: Support for other types than Float64, should be done in setup of
+        # pressure system
+        AMGX.pin_memory(nzval)
+        AMGX.upload!(s.matrix,
+            row_ptr,
+            colval,
+            nzval
+        )
+        amgx.data[:storage] = s
+        amgx.data[:block_size] = 1
+        amgx.data[:n] = n
+        amgx.data[:buffer_p] = AMGX.CUDA.CuVector{Tv}(undef, n)
+    else
+        pb = amgx.data[:buffer_p]
+        s = amgx.data[:storage]
+        A_gpu = s.matrix
+        @assert nnz(A) == nnz(A_gpu)
+        AMGX.replace_coefficients!(A_gpu, A.At.nzval)
+    end
+    AMGX.setup!(s.solver, s.matrix)
+    return amgx
+end
+
+function JutulDarcy.gpu_cpr_setup_buffers!(cpr, J_bsr, r_cu, op)
+    # TODO: Rename this.
+    data = cpr.pressure_precond.data
+    cpr_s = cpr.storage
+    if !haskey(data, :w_p)
+        Tv = eltype(r_cu)
+        n = length(r_cu)
+        bz = cpr_s.block_size
+        cpr.pressure_precond.data[:buffer_full] = similar(r_cu)
+        bz_w, n_w = size(cpr_s.w_p)
+        @assert n == bz*n_w
+        data[:w_p] = AMGX.CUDA.CuMatrix{Tv}(undef, bz_w, n_w)
+        data[:main_system] = J_bsr
+        data[:operator] = op
+    end
+    # Put updated weights on GPU
+    w_p_cpu = cpr_s.w_p
+    w_p_gpu = data[:w_p]
+    @assert size(w_p_cpu) == size(w_p_gpu)
+    copyto!(w_p_gpu, w_p_cpu)
+    return cpr
+end

ext/JutulDarcyAMGXExt/precond.jl

@@ -0,0 +1,53 @@
+function JutulDarcy.gpu_update_preconditioner!(prec::JutulDarcy.AMGXPreconditioner, lsys, model, storage, recorder, executor, krylov, J_bsr, r_cu, op)
+    nzval = nonzeros(J_bsr)
+    bz = Jutul.block_size(lsys[1, 1])
+    if haskey(prec.data, :storage)
+        s = prec.data[:storage]
+        AMGX.replace_coefficients!(s.matrix, nzval)
+    else
+        Tv = eltype(J_bsr)
+        if Tv == Float64
+            amgx_mode = AMGX.dDDI
+        else
+            amgx_mode = AMGX.dFFI
+        end
+        n, m = size(J_bsr)
+        @assert n == m
+        N = n ÷ bz
+        config = AMGX.Config(prec.settings)
+        s = AMGXStorage(config, amgx_mode)
+
+        AMGX.set_zero!(s.x, N, block_dim = bz)
+        AMGX.set_zero!(s.r, N, block_dim = bz)
+
+        row_ptr = J_bsr.rowPtr
+        colval = J_bsr.colVal
+        AMGX.upload!(
+            s.matrix,
+            row_ptr,
+            colval,
+            nzval,
+            block_dims = (bz, bz)
+        )
+        prec.data[:storage] = s
+        prec.data[:n] = N
+        prec.data[:block_size] = bz
+    end
+    AMGX.setup!(s.solver, s.matrix)
+end
+
+function Jutul.operator_nrows(prec::JutulDarcy.AMGXPreconditioner)
+    return prec.data[:n]*prec.data[:block_size]
+end
+
+function Jutul.apply!(x, amgx::JutulDarcy.AMGXPreconditioner, r)
+    # TODO: This sync call is maybe needd?
+    AMGX.CUDA.synchronize()
+    s = amgx.data[:storage]
+    bz = amgx.data[:block_size]::Int
+    n = amgx.data[:n]::Int
+    AMGX.upload!(s.r, r, block_dim = bz)
+    AMGX.set_zero!(s.x, n, block_dim = bz)
+    AMGX.solve!(s.x, s.solver, s.r)
+    AMGX.download!(x, s.x)
+end

ext/JutulDarcyCUDAExt/JutulDarcyCUDAExt.jl

@@ -0,0 +1,11 @@
+module JutulDarcyCUDAExt
+    using Jutul, JutulDarcy, CUDA, LinearAlgebra, SparseArrays
+    import Jutul: @tic
+
+    timeit_debug_enabled() = Jutul.timeit_debug_enabled()
+
+    include("ilu0.jl")
+    include("krylov.jl")
+    include("cuda_utils.jl")
+    include("cpr.jl")
+end

ext/JutulDarcyCUDAExt/cpr.jl

@@ -0,0 +1,65 @@
+function JutulDarcy.gpu_reduce_residual!(r_p, w_p, r)
+    @assert eltype(r_p) == eltype(w_p) == eltype(r)
+    @assert length(w_p) == length(r)
+    @assert size(w_p, 2) == length(r_p)
+    kernel = CUDA.@cuda launch = false reduce_residual_kernel!(r_p, w_p, r)
+    threads, blocks = kernel_helper(kernel, r_p)
+    CUDA.@sync begin
+        kernel(r_p, w_p, r; threads, blocks)
+    end
+    return r_p
+end
+
+function reduce_residual_kernel!(r_p, w_p, r_ps)
+    index = threadIdx().x
+    stride = blockDim().x
+    T = eltype(r_p)
+    ncomp = size(w_p, 1)
+    for cell in index:stride:length(r_p)
+        v = zero(T)
+        for comp in 1:ncomp
+            ix = (cell-1)*ncomp + comp
+            @inbounds wi = w_p[comp, cell]
+            @inbounds ri = r_ps[ix]
+            v += wi*ri
+        end
+        @inbounds r_p[cell] = v
+    end
+    return nothing
+end
+
+function JutulDarcy.gpu_increment_pressure!(x, dp)
+    n = length(x)
+    bz = div(n, length(dp))
+    @assert bz > 0
+    if false
+        kernel = CUDA.@cuda launch=false gpu_increment_pressure_kernel!(x, dp, bz)
+        threads, blocks = kernel_helper(kernel, dp)
+        threads = blocks = 1
+
+        CUDA.@sync begin
+        kernel(x, dp, bz; threads, blocks)
+        end
+    else
+        x_view = view(x, 1:bz:(n+1-bz))
+        @. x_view += dp
+    end
+    return x
+end
+
+function gpu_increment_pressure_kernel!(x, dp, bz)
+    index = threadIdx().x
+    stride = blockDim().x
+    for cell in index:stride:length(dp)
+        x[(cell-1)*bz + 1] += dp[cell]
+    end
+    return nothing
+end
+
+function kernel_helper(kernel, V)
+    config = launch_configuration(kernel.fun)
+    threads = min(length(V), config.threads)
+    blocks = cld(length(V), threads)
+
+    return (threads, blocks)
+end