Rework: Extensions

Project.toml

@@ -6,11 +6,11 @@ version = "0.1.2"
 [deps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4"
 GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0"
+JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
@@ -21,6 +21,12 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+JustRelaxCUDAExt = "CUDA"
+
 [compat]
 AMDGPU = "0.6, 0.7, 0.8"
 Adapt = "3"

ext/JustRelaxCUDAExt.jl

@@ -0,0 +1,179 @@
+module JustRelaxCUDAExt
+
+using CUDA
+using JustRelax: JustRelax
+import JustRelax: PTArray
+
+JustRelax.PTArray(::Type{CUDABackend}) = CuArray
+
+module JustRelax2D
+
+    using JustRelax: JustRelax
+    using CUDA
+    using StaticArrays
+    using CellArrays
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    using ImplicitGlobalGrid
+    using GeoParams, LinearAlgebra, Printf
+    using MPI
+
+    import JustRelax:
+        IGG,
+        BackendTrait,
+        CPUBackendTrait,
+        backend,
+        CPUBackend,
+        Geometry,
+        @cell
+
+    @init_parallel_stencil(CUDA, Float64, 2)
+
+    include("../src/common.jl")
+    include("../src/stokes/Stokes2D.jl")
+
+    # add CUDA traits
+    struct CUDABackendTrait <: BackendTrait end
+
+    @inline backend(::CuArray) = CUDABackendTrait()
+    @inline backend(::Type{<:CuArray}) = CUDABackendTrait()
+
+    # Types
+    function JustRelax.JustRelax2D.StokesArrays(
+        ::Type{CUDABackend}, ni::Vararg{Integer,N}
+    ) where {N}
+        return StokesArrays(tuple(ni...))
+    end
+
+    function JustRelax.JustRelax2D.StokesArrays(
+        ::Type{CUDABackend}, ni::NTuple{N,Integer}
+    ) where {N}
+        return StokesArrays(ni)
+    end
+
+    function JustRelax.JustRelax2D.ThermalArrays(
+        ::Type{CUDABackend}, ni::NTuple{N,Number}
+    ) where {N}
+        return ThermalArrays(ni...)
+    end
+
+    function JustRelax.JustRelax2D.ThermalArrays(
+        ::Type{CUDABackend}, ni::Vararg{Number,N}
+    ) where {N}
+        return ThermalArrays(ni...)
+    end
+
+    function JustRelax.JustRelax2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases)
+        return PhaseRatio(ni, num_phases)
+    end
+
+    # Boundary conditions
+    function JustRelax.JustRelax2D.flow_bcs!(
+        ::CUDABackendTrait, stokes::StokesArrays, bcs
+    )
+        return _flow_bcs!(bcs, @velocity(stokes))
+    end
+
+    function flow_bcs!(
+        ::CUDABackendTrait, stokes::StokesArrays, bcs
+    )
+        return _flow_bcs!(bcs, @velocity(stokes))
+    end
+
+    function JustRelax.JustRelax2D.thermal_bcs!(
+        ::CUDABackendTrait, thermal::ThermalArrays, bcs
+    )
+        return thermal_bcs!(thermal.T, bcs)
+    end
+
+    function thermal_bcs!(
+        ::CUDABackendTrait, thermal::ThermalArrays, bcs
+    )
+        return thermal_bcs!(thermal.T, bcs)
+    end
+
+    # Phases
+    function JustRelax.JustRelax2D.phase_ratios_center(
+        ::CUDABackendTrait, phase_ratios::PhaseRatio, particles, grid::Geometry, phases
+    )
+        return _phase_ratios_center(phase_ratios, particles, grid, phases)
+    end
+
+    # Rheology
+    ## viscosity
+    function JustRelax.JustRelax2D.compute_viscosity!(
+        ::CUDABackendTrait, stokes, ν, args, rheology, cutoff
+    )
+        return _compute_viscosity!(stokes, ν, args, rheology, cutoff)
+    end
+    function JustRelax.JustRelax2D.compute_viscosity!(
+        ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff
+    )
+        return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff)
+    end
+    function JustRelax.JustRelax2D.compute_viscosity!(
+        η, ν, εII::CuArray, args, rheology, cutoff
+    )
+        return compute_viscosity!(η, ν, εII, args, rheology, cutoff)
+    end
+
+    function compute_viscosity!(
+        ::CUDABackendTrait, stokes, ν, args, rheology, cutoff
+    )
+        return _compute_viscosity!(stokes, ν, args, rheology, cutoff)
+    end
+    function compute_viscosity!(
+        ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff
+    )
+        return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff)
+    end
+    function compute_viscosity!(
+        η, ν, εII::CuArray, args, rheology, cutoff
+    )
+        return compute_viscosity!(η, ν, εII, args, rheology, cutoff)
+    end
+
+    ## Stress
+    JustRelax.JustRelax2D.tensor_invariant!(A::SymmetricTensor) = tensor_invariant!(A)
+
+    ## Buoyancy forces
+    function JustRelax.JustRelax2D.compute_ρg!(ρg::CuArray, rheology, args)
+        return compute_ρg!(ρg, rheology, args)
+    end
+    function JustRelax.JustRelax2D.compute_ρg!(
+        ρg::CuArray, phase_ratios::PhaseRatio, rheology, args
+    )
+        return compute_ρg!(ρg, phase_ratios, rheology, args)
+    end
+
+    # Interpolations
+    function JustRelax.JustRelax2D.temperature2center!(
+        ::CUDABackendTrait, thermal::ThermalArrays
+    )
+        return _temperature2center!(thermal)
+    end
+    function JustRelax.JustRelax2D.vertex2center!(center::T, vertex::T) where {T<:CuArray}
+        return vertex2center!(center, vertex)
+    end
+    function JustRelax.JustRelax2D.center2vertex!(vertex::T, center::T) where {T<:CuArray}
+        return center2vertex!(vertex, center)
+    end
+
+    function JustRelax.JustRelax2D.center2vertex!(
+        vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T
+    ) where {T<:CuArray}
+        return center2vertex!(
+            vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy
+        )
+    end
+
+    # Solvers
+    JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) 
+
+    # Utils
+    JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) = compute_dt(S, di, dt_diff, I::IGG)
+    JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff) = compute_dt(S, di, dt_diff)
+    JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di) = compute_dt(S, di)
+
+end
+
+end

miniapps/benchmarks/stokes2D/RunStokesBench2D.jl

@@ -1,14 +1,10 @@
-using JustRelax, Printf, LinearAlgebra
+using JustRelax, JustRelax.JustRelax2D, Printf, LinearAlgebra
 using MPI: MPI
 using GLMakie
 
 using ParallelStencil
 @init_parallel_stencil(Threads, Float64, 2)
 
-# setup ParallelStencil.jl environment
-model = PS_Setup(:cpu, Float64, 2) # :cpu | :CUDA | :AMDGPU
-environment!(model)
-
 # choose benchmark
 benchmark = :solcx
 

miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl

@@ -50,7 +50,7 @@ function elastic_buildup(;
 
     ## Allocate arrays needed for every Stokes problem
     # general stokes arrays
-    stokes    = StokesArrays(ni, ViscoElastic)
+    stokes    = StokesArrays(ni)
     # general numerical coeffs for PT stokes
     pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL=1 / √2.1)
 
@@ -78,15 +78,16 @@ function elastic_buildup(;
             di,
             flow_bcs,
             ρg,
-            η,
             Gc,
             Kb,
             dt,
             igg;
-            iterMax=150e3,
-            nout=1000,
-            b_width=(4, 4, 1),
-            verbose=true,
+            kwargs = (;
+                iterMax=150e3,
+                nout=1000,
+                b_width=(4, 4, 1),
+                verbose=true,
+            )
         )
 
         @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ))

miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl

@@ -1,14 +1,8 @@
-using JustRelax, JustRelax.DataIO
-import JustRelax.@cell
-
-# setup ParallelStencil.jl environment
-model = PS_Setup(:Threads, Float64, 2)
-environment!(model)
-
+using JustRelax, JustRelax.JustRelax2D
 using JustPIC, JustPIC._2D
 const backend = CPUBackend
 
-using ParallelStencil
+using ParallelStencil, ParallelStencil.FiniteDifferences2D
 @init_parallel_stencil(Threads, Float64, 2)
 
 # Load script dependencies
@@ -44,10 +38,10 @@ function init_phases!(phases, particles)
 
         @inbounds for ip in JustRelax.cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+            @cell(index[ip, i, j]) == 0 && continue
 
-            x = JustRelax.@cell px[ip, i, j]
-            depth = -(JustRelax.@cell py[ip, i, j])
+            x = @cell px[ip, i, j]
+            depth = -(@cell py[ip, i, j]) 
             @cell phases[ip, i, j] = 2.0
 
             if 0e0 ≤ depth ≤ 100e3
@@ -124,7 +118,7 @@ function main(igg, nx, ny)
     # Elliptical temperature anomaly
     init_phases!(pPhases, particles)
     phase_ratios  = PhaseRatio(ni, length(rheology))
-    @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases)
+    phase_ratios_center(phase_ratios, particles, grid, pPhases)
     # ----------------------------------------------------
 
     # STOKES ---------------------------------------------
@@ -139,14 +133,10 @@ function main(igg, nx, ny)
 
     # Buoyancy forces & rheology
     ρg               = @zeros(ni...), @zeros(ni...)
-    η                = @ones(ni...)
     args             = (; T = thermal.Tc, P = stokes.P, dt = Inf)
-    @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P))
+    compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P))
     @parallel init_P!(stokes.P, ρg[2], xci[2])
-    @parallel (@idx ni) compute_viscosity!(
-        η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf)
-    )
-    η_vep            = copy(η)
+    compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf))
 
     # Boundary conditions
     flow_bcs         = FlowBoundaryConditions(;
@@ -160,7 +150,7 @@ function main(igg, nx, ny)
         ax1 = Axis(fig[1,1], aspect = 2/3, title = "T")
         ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)")
         scatter!(ax1, Array(ρg[2][:]./9.81), Y./1e3)
-        scatter!(ax2, Array(log10.(η[:])), Y./1e3)
+        scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3)
         # scatter!(ax2, Array(stokes.P[:]), Y./1e3)
         ylims!(ax1, minimum(xvi[2])./1e3, 0)
         ylims!(ax2, minimum(xvi[2])./1e3, 0)
@@ -181,29 +171,25 @@ function main(igg, nx, ny)
 
         # Stokes solver ----------------
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
-
-        @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P))
-        @parallel init_P!(stokes.P, ρg[2], xci[2])
-        @parallel (@idx ni) compute_viscosity!(
-            η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf)
-        )
+        compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P))
+        compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf))
 
         solve!(
             stokes,
             pt_stokes,
             di,
             flow_bcs,
             ρg,
-            η,
-            η_vep,
             phase_ratios,
             rheology,
             args,
             dt,
             igg;
-            iterMax = 150e3,
-            nout=1e3,
-            viscosity_cutoff=(-Inf, Inf)
+            kwargs = (;
+                iterMax = 150e3,
+                nout=1e3,
+                viscosity_cutoff=(-Inf, Inf)
+            )
         )
         dt = compute_dt(stokes, di) / 2
         # ------------------------------
@@ -217,8 +203,8 @@ function main(igg, nx, ny)
         inject = check_injection(particles)
         inject && inject_particles_phase!(particles, pPhases, (), (), xvi)
         # update phase ratios
-        @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases)
-
+        phase_ratios_center(phase_ratios, particles, grid, pPhases)
+ 
         @show it += 1
         t        += dt
 
@@ -227,7 +213,7 @@ function main(igg, nx, ny)
             nt = 5
             fig = Figure(resolution = (900, 900), title = "t = $t")
             ax  = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs")
-            heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)), colormap = :grayC)
+            heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC)
             arrows!(
                 ax,
                 xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))...,

miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl

@@ -1,6 +1,5 @@
 # using CUDA
 using JustRelax, JustRelax.DataIO
-import JustRelax.@cell
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:Threads, Float64, 2)
@@ -43,10 +42,10 @@ function init_phases!(phases, particles, A)
 
         @inbounds for ip in JustRelax.cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+            @cell(index[ip, i, j]) == 0 && continue
 
-            x = JustRelax.@cell px[ip, i, j]
-            depth = -(JustRelax.@cell py[ip, i, j]) 
+            x = @cell px[ip, i, j]
+            depth = -(@cell py[ip, i, j]) 
             @cell phases[ip, i, j] = 2.0
 
             if 0e0 ≤ depth ≤ 100e3