Merge pull request #2825 from CliMA/ncc/drop-ifhasamgx

Use AMGX v0.1.3 + drop `@ifhasamgx`

Manifest.toml

@@ -2,13 +2,13 @@
 
 julia_version = "1.8.2"
 manifest_format = "2.0"
-project_hash = "a4995b60686d82011ed52a33a7533a8246769dc0"
+project_hash = "f6464a834537dc16cda28b2f9dd3b7da0b24f89d"
 
 [[deps.AMGX]]
 deps = ["AMGX_jll", "CEnum", "CUDA", "JSON", "Libdl", "SparseArrays"]
-git-tree-sha1 = "f31d1873273b84f6694f49d35c3a894a15e21e8a"
+git-tree-sha1 = "c1686ca92a2ceb2031fbfe0860a86f4c71bd1ac8"
 uuid = "c963dde9-0319-47f5-bf0c-b07d3c80ffa6"
-version = "0.1.2"
+version = "0.1.3"
 
 [[deps.AMGX_jll]]
 deps = ["Artifacts", "CUDA_jll", "JLLWrappers", "Libdl", "Pkg"]

Project.toml

@@ -39,9 +39,9 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
-AMGX = "0.1"
 Adapt = "^3"
 AlgebraicMultigrid = "0.5"
+AMGX = "0.1.3"
 CUDA = "3.8, 3.9"
 CUDAKernels = "0.3.3, 0.4"
 Crayons = "^4"

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -8,7 +8,6 @@ export
 using KernelAbstractions: @index, @kernel, Event, MultiEvent, NoneEvent
 using KernelAbstractions.Extras.LoopInfo: @unroll
 
-using Oceananigans: @ifhasamgx
 using Oceananigans.Utils
 using Oceananigans.Utils: launch!
 

src/Models/HydrostaticFreeSurfaceModels/matrix_implicit_free_surface_solver.jl

@@ -141,4 +141,4 @@ end
     end
 end
 
-@ifhasamgx finalize_solver!(s::MatrixImplicitFreeSurfaceSolver) = finalize_solver!(s.matrix_iterative_solver)
+finalize_solver!(s::MatrixImplicitFreeSurfaceSolver) = finalize_solver!(s.matrix_iterative_solver)

src/Models/HydrostaticFreeSurfaceModels/mg_implicit_free_surface_solver.jl

@@ -10,7 +10,7 @@ using Oceananigans.Solvers: constructors, arch_sparse_matrix, update_diag!, unpa
 using Oceananigans.Utils: prettysummary
 using SparseArrays: _insert!
 using CUDA.CUSPARSE: CuSparseMatrixCSR
-@ifhasamgx using AMGX
+using AMGX
 
 import Oceananigans.Solvers: solve!, finalize_solver!
 import Oceananigans.Architectures: architecture
@@ -165,7 +165,7 @@ function solve!(η, implicit_free_surface_solver::MGImplicitFreeSurfaceSolver{CP
     return nothing
 end
 
-@ifhasamgx function solve!(η, implicit_free_surface_solver::MGImplicitFreeSurfaceSolver{GPU}, rhs, g, Δt)
+function solve!(η, implicit_free_surface_solver::MGImplicitFreeSurfaceSolver{GPU}, rhs, g, Δt)
     solver = implicit_free_surface_solver.multigrid_solver
 
     # if `Δt` changed then re-compute the matrix elements

src/Oceananigans.jl

@@ -139,23 +139,12 @@ import Base:
     push!
 
 
-# TODO: find a way to check whether the libraries for AMGX and NETCDF 
-# (libamgxsh and libnetcdf, respectively) are installed on the machine
-"Boolean denoting whether AMGX.jl can be loaded on machine."
-const hasamgx   = @static (Sys.islinux() && Sys.ARCH == :x86_64) ? true : false
+# TODO: find a way to check whether the libraries for NETCDF 
+# (libnetcdf) are installed on the machine
 
 "Boolean denoting whether NCDatasets.jl can be loaded on machine."
 const hasnetcdf = @static (Sys.islinux() && Sys.ARCH == :x86_64) ? true : false
 
-"""
-    @ifhasamgx expr
-
-Evaluate `expr` only if `hasamgx == true`.
-"""
-macro ifhasamgx(expr)
-    hasamgx ? :($(esc(expr))) : :(nothing) 
-end
-
 """
     @ifnetcdf expr
 

src/Solvers/Solvers.jl

@@ -20,7 +20,6 @@ using Oceananigans.Grids
 using Oceananigans.BoundaryConditions
 using Oceananigans.Fields
 
-using Oceananigans: @ifhasamgx, hasamgx
 using Oceananigans.Grids: unpack_grid
 
 """

src/Solvers/heptadiagonal_iterative_solver.jl

@@ -324,7 +324,7 @@ function solve!(x, solver::HeptadiagonalIterativeSolver, b, Δt)
     return nothing
 end
 
-@ifhasamgx finalize_solver!(p::HeptadiagonalIterativeSolver) = finalize_solver!(p.preconditioner)
+finalize_solver!(p::HeptadiagonalIterativeSolver) = finalize_solver!(p.preconditioner)
 
 function Base.show(io::IO, solver::HeptadiagonalIterativeSolver)
     print(io, "Matrix-based iterative solver with: \n")

src/Solvers/multigrid_solver.jl

@@ -2,10 +2,13 @@ using Oceananigans.Architectures: architecture
 using LinearAlgebra
 using AlgebraicMultigrid: _solve!, RugeStubenAMG, ruge_stuben, SmoothedAggregationAMG, smoothed_aggregation
 using CUDA
-@ifhasamgx using AMGX
+using AMGX
 
 import Oceananigans.Architectures: architecture
 
+"Boolean denoting whether AMGX.jl can be loaded on machine."
+const hasamgx   = @static (Sys.islinux() && Sys.ARCH == :x86_64) ? true : false
+
 abstract type MultigridSolver{A, G, L, T, F} end
 
 mutable struct MultigridCPUSolver{A, G, L, T, F, R, M} <: MultigridSolver{A, G, L, T, F}
@@ -154,16 +157,16 @@ function MultigridSolver_on_architecture(::CPU;
                               )
 end
 
-@ifhasamgx function MultigridSolver_on_architecture(::GPU;
-                                                    template_field::AbstractField,
-                                                    maxiter,
-                                                    reltol,
-                                                    abstol,
-                                                    amg_algorithm,
-                                                    matrix,
-                                                    x_array,
-                                                    b_array
-                                                    )
+function MultigridSolver_on_architecture(::GPU;
+                                         template_field::AbstractField,
+                                         maxiter,
+                                         reltol,
+                                         abstol,
+                                         amg_algorithm,
+                                         matrix,
+                                         x_array,
+                                         b_array
+                                         )
 
     amgx_solver = AMGXMultigridSolver(matrix, maxiter, reltol, abstol)
 
@@ -179,7 +182,7 @@ end
                               )
 end
 
-@ifhasamgx function AMGXMultigridSolver(matrix::CuSparseMatrixCSC, maxiter = 1, reltol = sqrt(eps(eltype(matrix))), abstol = 0)
+function AMGXMultigridSolver(matrix::CuSparseMatrixCSC, maxiter = 1, reltol = sqrt(eps(eltype(matrix))), abstol = 0)
     tolerance, convergence = reltol == 0 ? (abstol, "ABSOLUTE") : (reltol, "RELATIVE_INI_CORE")
     try
         global config = AMGX.Config(Dict("monitor_residual" => 1, "max_iters" => maxiter, "store_res_history" => 1, "tolerance" => tolerance, "convergence" => convergence))
@@ -293,9 +296,11 @@ function solve!(x, solver::MultigridCPUSolver, b; kwargs...)
             kwargs...)
 
     interior(x) .= reshape(solver.x_array, Nx, Ny, Nz)
+
+    return nothing
 end
 
-@ifhasamgx function solve!(x, solver::MultigridGPUSolver, b; kwargs...)
+function solve!(x, solver::MultigridGPUSolver, b; kwargs...)
     Nx, Ny, Nz = size(b)
 
     solver.b_array .= reshape(interior(b), Nx * Ny * Nz)
@@ -308,6 +313,8 @@ end
     AMGX.copy!(solver.x_array, s.device_x)
 
     interior(x) .= reshape(solver.x_array, Nx, Ny, Nz)
+
+    return nothing
 end
 
 """
@@ -318,7 +325,7 @@ This function needs to be called before creating a multigrid solver on GPU.
 """
 function initialize_AMGX(::GPU)
     try
-        @ifhasamgx AMGX.initialize(); AMGX.initialize_plugins()
+        AMGX.initialize(); AMGX.initialize_plugins()
     catch e
         @info "It appears AMGX was not finalized. Have you called `finalize_AMGX`?"
         AMGX.finalize_plugins()
@@ -337,13 +344,13 @@ Finalize the AMGX package required to use the multigrid solver on `architecture`
 This should be called after `finalize_solver!`.
 """
 function finalize_AMGX(::GPU)
-    @ifhasamgx AMGX.finalize_plugins(); AMGX.finalize()
+    AMGX.finalize_plugins(); AMGX.finalize()
 end
 
 finalize_AMGX(::CPU) = nothing
 
 
-@ifhasamgx function finalize_solver!(s::AMGXMultigridSolver)
+function finalize_solver!(s::AMGXMultigridSolver)
     @info "Finalizing the AMGX Multigrid solver on GPU"
     close(s.device_matrix)
     close(s.device_x)
@@ -355,7 +362,7 @@ finalize_AMGX(::CPU) = nothing
     return nothing
 end
 
-@ifhasamgx finalize_solver!(solver::MultigridGPUSolver) = finalize_solver!(solver.amgx_solver)
+finalize_solver!(solver::MultigridGPUSolver) = finalize_solver!(solver.amgx_solver)
 
 finalize_solver!(::MultigridCPUSolver) = nothing
 

src/Solvers/sparse_preconditioners.jl

@@ -182,28 +182,26 @@ end
 
 multigrid_preconditioner(A::AbstractMatrix) = aspreconditioner(create_multilevel(RugeStubenAMG(), A))
 
-@ifhasamgx multigrid_preconditioner(A::CuSparseMatrixCSC) = MultigridGPUPreconditioner(AMGXMultigridSolver(A))
+multigrid_preconditioner(A::CuSparseMatrixCSC) = MultigridGPUPreconditioner(AMGXMultigridSolver(A))
 
 struct MultigridGPUPreconditioner{AMGXMultigridSolver}
     amgx_solver :: AMGXMultigridSolver
 end
 
 import LinearAlgebra: \, *, ldiv!, mul!
 
-@ifhasamgx begin
-    ldiv!(p::MultigridGPUPreconditioner, b) = copyto!(b, p \ b)
+ldiv!(p::MultigridGPUPreconditioner, b) = copyto!(b, p \ b)
 
-    function ldiv!(x, p::MultigridGPUPreconditioner, b)
-        x .= 0
-        AMGX.upload!(p.amgx_solver.device_b, b)
-        AMGX.upload!(p.amgx_solver.device_x, x)
-        AMGX.solve!(p.amgx_solver.device_x, p.amgx_solver.solver, p.amgx_solver.device_b)
-        AMGX.copy!(x, p.amgx_solver.device_x)
-    end
+function ldiv!(x, p::MultigridGPUPreconditioner, b)
+    x .= 0
+    AMGX.upload!(p.amgx_solver.device_b, b)
+    AMGX.upload!(p.amgx_solver.device_x, x)
+    AMGX.solve!(p.amgx_solver.device_x, p.amgx_solver.solver, p.amgx_solver.device_b)
+    AMGX.copy!(x, p.amgx_solver.device_x)
+end
 
-    mul!(b, p::MultigridGPUPreconditioner, x) = mul!(b, p.amgx_solver.csr_matrix, x)
+mul!(b, p::MultigridGPUPreconditioner, x) = mul!(b, p.amgx_solver.csr_matrix, x)
 
-    \(p::MultigridGPUPreconditioner, b) = ldiv!(similar(b), p, b)
-
-    finalize_solver!(p::MultigridGPUPreconditioner) = finalize_solver!(p.amgx_solver)
-end
+\(p::MultigridGPUPreconditioner, b) = ldiv!(similar(b), p, b)
+
+finalize_solver!(p::MultigridGPUPreconditioner) = finalize_solver!(p.amgx_solver)