add GPU compatibility

PTsolvers · Oct 22, 2024 · 1afc888 · 1afc888
1 parent f1d2530
commit 1afc888
Show file tree

Hide file tree

Showing 8 changed files with 136 additions and 23 deletions.
diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -55,6 +55,10 @@ function JR2D.WENO5(
     return WENO5(method, tuple(ni...))
 end
 
+function JR2D.RockRatio(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N}
+    return RockRatio(ni...)
+end
+
 function JR2D.PTThermalCoeffs(
     ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3
 )
@@ -286,6 +290,10 @@ function JR2D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs)
     return _solve!(stokes, args...; kwargs...)
 end
 
+function JR2D.solve_VariationalStokes!(::AMDGPUBackendTrait, stokes, args...; kwargs)
+    return _solve_VS!(stokes, args...; kwargs...)
+end
+
 function JR2D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs)
     return _heatdiffusion_PT!(thermal, args...; kwargs...)
 end
@@ -391,4 +399,13 @@ function JR2D.rotate_stress_particles!(
     return nothing
 end
 
+# rock ratios
+
+function JR2D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase
+) where {T,nD,N}
+    update_rock_ratio!(ϕ, phase_ratios, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -55,6 +55,10 @@ function JR3D.WENO5(
     return WENO5(method, tuple(ni...))
 end
 
+function JR3D.RockRatio(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N}
+    return RockRatio(ni...)
+end
+
 function JR3D.PTThermalCoeffs(
     ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3
 )
@@ -285,11 +289,16 @@ end
 function displacement2velocity!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, dt)
     return _displacement2velocity!(stokes, dt)
 end
+
 # Solvers
 function JR3D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs)
     return _solve!(stokes, args...; kwargs...)
 end
 
+function JR3D.solve_VariationalStokes!(::AMDGPUBackendTrait, stokes, args...; kwargs)
+    return _solve_VS!(stokes, args...; kwargs...)
+end
+
 function JR3D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs)
     return _heatdiffusion_PT!(thermal, args...; kwargs...)
 end
@@ -407,4 +416,13 @@ function JR3D.rotate_stress_particles!(
     return nothing
 end
 
+# rock ratios
+
+function JR3D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{ROCArray{T,nD},N}, phase_ratios, air_phase
+) where {T,nD,N}
+    update_rock_ratio!(ϕ, phase_ratios, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -30,6 +30,7 @@ include("../../common.jl")
 include("../../stokes/Stokes2D.jl")
 
 # Types
+
 function JR2D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N}
     return StokesArrays(ni)
 end
@@ -46,6 +47,14 @@ function JR2D.WENO5(::Type{CUDABackend}, method::Val{T}, ni::NTuple{N,Integer})
     return WENO5(method, tuple(ni...))
 end
 
+function JR2D.RockRatio(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N}
+    return RockRatio(ni...)
+end
+
+function JR2D.RockRatio(::Type{CUDABackend}, ni::Vararg{Integer,N}) where {N}
+    return RockRatio(ni...)
+end
+
 function JR2D.PTThermalCoeffs(
     ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3
 )
@@ -270,6 +279,10 @@ function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs)
     return _solve!(stokes, args...; kwargs...)
 end
 
+function JR2D.solve_VariationalStokes!(::CUDABackendTrait, stokes, args...; kwargs)
+    return _solve_VS!(stokes, args...; kwargs...)
+end
+
 function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs)
     return _heatdiffusion_PT!(thermal, args...; kwargs...)
 end
@@ -372,4 +385,13 @@ function JR2D.rotate_stress_particles!(
     return nothing
 end
 
+# rock ratios
+
+function JR2D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -45,6 +45,10 @@ function JR3D.WENO5(::Type{CUDABackend}, method::Val{T}, ni::NTuple{N,Integer})
     return WENO5(method, tuple(ni...))
 end
 
+function JR3D.RockRatio(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N}
+    return RockRatio(ni...)
+end
+
 function JR3D.PTThermalCoeffs(
     ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3
 )
@@ -289,6 +293,10 @@ function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs)
     return _solve!(stokes, args...; kwargs...)
 end
 
+function JR3D.solve_VariationalStokes!(::CUDABackendTrait, stokes, args...; kwargs)
+    return _solve_VS!(stokes, args...; kwargs...)
+end
+
 function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs)
     return _heatdiffusion_PT!(thermal, args...; kwargs...)
 end
@@ -404,4 +412,13 @@ function JR3D.rotate_stress_particles!(
     return nothing
 end
 
+# rock ratios
+
+function JR3D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, air_phase
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl
@@ -90,7 +90,7 @@ function _solve_VS!(
             compute_maxloc!(ητ, η; window=(1, 1))
             update_halo!(ητ)
 
-            @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes), ϕ, _di)
+            @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., ϕ, _di...)
 
             compute_P!(
                 θ,

diff --git a/src/variational_stokes/VelocityKernels.jl b/src/variational_stokes/VelocityKernels.jl
@@ -1,20 +1,54 @@
 @parallel_indices (I...) function compute_∇V!(
-    ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N}
-) where {T,N}
-    @inline d_xi(A) = _d_xi(A, _di[1], I...)
-    @inline d_yi(A) = _d_yi(A, _di[2], I...)
-    @inline d_zi(A) = _d_zi(A, _di[3], I...)
+    ∇V::AbstractArray{T,2}, Vx, Vy, ϕ::JustRelax.RockRatio, _dx, _dy
+) where {T}
+    @inline d_xi(A) = _d_xi(A, _dx, I...)
+    @inline d_yi(A) = _d_yi(A, _dy, I...)
 
-    f = d_xi, d_yi, d_zi
+    if isvalid_c(ϕ, I...)
+        @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy)
+    else
+        @inbounds ∇V[I...] = zero(T)
+    end
+    return nothing
+end
+
+@parallel_indices (I...) function compute_∇V!(
+    ∇V::AbstractArray{T,2}, Vx, Vy, Vz, ϕ::JustRelax.RockRatio, _dx, _dy, _dz
+) where {T}
+    @inline d_xi(A) = _d_xi(A, _dx, I...)
+    @inline d_yi(A) = _d_yi(A, _dy, I...)
+    @inline d_zi(A) = _d_zi(A, _dz, I...)
 
     if isvalid_c(ϕ, I...)
-        @inbounds ∇V[I...] = sum(f[i](V[i]) for i in 1:N)
+        @inbounds ∇V[I...] = d_xi(Vx) + d_yi(Vy) + d_zi(Vz)
     else
         @inbounds ∇V[I...] = zero(T)
     end
     return nothing
 end
 
+# @parallel_indices (I...) function compute_∇V!(
+#     ∇V::AbstractArray{T,N}, V::NTuple{N}, ϕ::JustRelax.RockRatio, _di::NTuple{N}
+# ) where {T,N}
+#     # @inline d_xi(A) = _d_xi(A, _di, I...)
+#     # @inline d_yi(A) = _d_yi(A, _di, I...)
+#     # @inline d_zi(A) = _d_zi(A, _di, I...)
+
+#     # f = d_xi, d_yi, d_zi
+#     f = _d_xi, _d_yi, _d_zi
+
+#     if isvalid_c(ϕ, I...)
+#         v = zero(T)
+#         for i in 1:N
+#             v += f[i](V[i], _di[I], I...)
+#         end
+#         @inbounds ∇V[I...] = v
+#     else
+#         @inbounds ∇V[I...] = zero(T)
+#     end
+#     return nothing
+# end
+
 @parallel_indices (i, j) function compute_V!(
     Vx::AbstractArray{T,2},
     Vy,

diff --git a/src/variational_stokes/mask.jl b/src/variational_stokes/mask.jl
@@ -1,12 +1,15 @@
+function RockRatio(::Type{CPUBackend}, ni::NTuple{N,Integer}) where {N}
+    return RockRatio(ni...)
+end
 
 function RockRatio(nx, ny)
     ni = nx, ny
     center = @zeros(ni...)
     vertex = @zeros(ni .+ 1...)
     Vx = @zeros(nx+1, ny) # no ghost nodes!
     Vy = @zeros(nx, ny+1) # no ghost nodes!
-
-    return JustRelax.RockRatio(center, vertex, Vx, Vy, nothing, nothing, nothing, nothing)
+    dummy = @zeros(1, 1) # because it cant be a Union{T, Nothing} type on the GPU....
+    return JustRelax.RockRatio(center, vertex, Vx, Vy, dummy, dummy, dummy, dummy)
 end
 
 function RockRatio(nx, ny, nz)
@@ -54,6 +57,7 @@ end
 @inline compute_rock_ratio(
     phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}
 ) where {N} = 1 - @index phase_ratio[air_phase, I...]
+
 @inline compute_air_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = @index phase_ratio[
     air_phase, I...
 ]

diff --git a/src/variational_stokes/types.jl b/src/variational_stokes/types.jl
@@ -5,25 +5,26 @@ struct RockRatio{T,N} <: AbstractMask
     vertex::T
     Vx::T
     Vy::T
-    Vz::Union{Nothing,T}
-    yz::Union{Nothing,T}
-    xz::Union{Nothing,T}
-    xy::Union{Nothing,T}
+    Vz::T
+    yz::T
+    xz::T
+    xy::T
 
     function RockRatio(
-        center::AbstractArray{F,N},
-        vertex::T,
-        Vx::T,
-        Vy::T,
-        Vz::Union{Nothing,T},
-        yz::Union{Nothing,T},
-        xz::Union{Nothing,T},
-        xy::Union{Nothing,T},
+        center::AbstractArray{F,N}, vertex::T, Vx::T, Vy::T, Vz::T, yz::T, xz::T, xy::T
     ) where {F,N,T}
         return new{T,N}(center, vertex, Vx, Vy, Vz, yz, xz, xy)
     end
 end
 
-RockRatio(ni::NTuple{N,Integer}) where {N} = RockRatio(ni...)
+RockRatio(::Type{CPUBackend}, ni::NTuple{N,Number}) where {N} = RockRatio(ni...)
+
+function RockRatio(::Number, ::Number)
+    throw(ArgumentError("RockRatio dimensions must be given as integers"))
+end
+
+function RockRatio(::Number, ::Number, ::Number)
+    throw(ArgumentError("RockRatio dimensions must be given as integers"))
+end
 
 Adapt.@adapt_structure RockRatio