diff --git a/Project.toml b/Project.toml index cec030eb..cc5b481e 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,7 @@ AMDGPU = "0.8, 0.9, 1" Adapt = "4" CUDA = "5" CellArrays = "0.2" -GeoParams = "0.6.4" +GeoParams = "0.6.7" HDF5 = "0.17.1" ImplicitGlobalGrid = "0.15" JLD2 = "0.4, 0.5" diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 5db87d39..3bab7627 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -215,10 +215,15 @@ function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTens end ## Buoyancy forces -function JR2D.compute_ρg!(ρg::ROCArray, rheology, args) +function JR2D.compute_ρg!(ρg::Union{ROCArray,NTuple{N,ROCArray}}, rheology, args) where {N} return compute_ρg!(ρg, rheology, args) end -function JR2D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) +function JR2D.compute_ρg!( + ρg::Union{ROCArray,NTuple{N,ROCArray}}, + phase_ratios::JustPIC.PhaseRatios, + rheology, + args, +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index b2e40b70..f9cb007b 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -215,15 +215,22 @@ function JR3D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTens end ## Buoyancy forces -function JR3D.compute_ρg!(ρg::ROCArray, rheology, args) +function JR3D.compute_ρg!(ρg::Union{ROCArray,NTuple{N,ROCArray}}, rheology, args) where {N} return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) +function JR3D.compute_ρg!( + ρg::Union{ROCArray,NTuple{N,ROCArray}}, + phase_ratios::JustPIC.PhaseRatios, + rheology, + args, +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end -function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios, rheology, args) +function JR3D.compute_ρg!( + ρg::Union{ROCArray,NTuple{N,ROCArray}}, phase_ratios, rheology, args +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index fd0f2a6f..d4eb3f27 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -202,11 +202,13 @@ function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor end ## Buoyancy forces -function JR2D.compute_ρg!(ρg::CuArray, rheology, args) +function JR2D.compute_ρg!(ρg::Union{CuArray,NTuple{N,CuArray}}, rheology, args) where {N} return compute_ρg!(ρg, rheology, args) end -function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) +function JR2D.compute_ρg!( + ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios::JustPIC.PhaseRatios, rheology, args +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 6cc0605b..34134f7d 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -219,13 +219,17 @@ function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor end ## Buoyancy forces -function JR3D.compute_ρg!(ρg::CuArray, rheology, args) +function JR3D.compute_ρg!(ρg::Union{CuArray,NTuple{N,CuArray}}, rheology, args) where {N} return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) +function JR3D.compute_ρg!( + ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios::JustPIC.PhaseRatios, rheology, args +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end -function JR3D.compute_ρg!(ρg::CuArray, phase_ratios, rheology, args) +function JR3D.compute_ρg!( + ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios, rheology, args +) where {N} return compute_ρg!(ρg, phase_ratios, rheology, args) end diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 946cab3b..126cdfb6 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -4,7 +4,10 @@ Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. """ function compute_ρg!(ρg, rheology, args) - ni = size(ρg) + _size(x::AbstractArray) = size(x) + _size(x::NTuple) = size(x[1]) + + ni = _size(ρg) @parallel (@idx ni) compute_ρg_kernel!(ρg, rheology, args) return nothing end @@ -15,6 +18,16 @@ end return nothing end +@parallel_indices (I...) function compute_ρg_kernel!( + ρg::NTuple{N,AbstractArray}, rheology, args +) where {N} + args_ijk = ntuple_idx(args, I...) + gᵢ = compute_gravity(first(rheology)) + ρgᵢ = compute_buoyancies(rheology, args_ijk, gᵢ, Val(N)) + fill_density!(ρg, ρgᵢ, I...) + return nothing +end + """ compute_ρg!(ρg, phase_ratios, rheology, args) @@ -22,7 +35,10 @@ Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object The `phase_ratios` are used to compute the density of the composite rheology. """ function compute_ρg!(ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args) - ni = size(ρg) + _size(x::AbstractArray) = size(x) + _size(x::NTuple) = size(x[1]) + + ni = _size(ρg) @parallel (@idx ni) compute_ρg_kernel!(ρg, phase_ratios.center, rheology, args) return nothing end @@ -33,7 +49,50 @@ end return nothing end +@parallel_indices (I...) function compute_ρg_kernel!( + ρg::NTuple{N,AbstractArray}, phase_ratios, rheology, args +) where {N} + args_ijk = ntuple_idx(args, I...) + gᵢ = compute_gravity(first(rheology)) + ρgᵢ = compute_buoyancies(rheology, @cell(phase_ratios[I...]), args_ijk, gᵢ, Val(N)) + fill_density!(ρg, ρgᵢ, I...) + return nothing +end + ## Inner buoyancy force kernels +@generated function fill_density!(ρg::NTuple{N}, ρgᵢ::NTuple{N}, I::Vararg{Int,N}) where {N} + quote + Base.@nexprs $N i -> ρg[i][I...] = ρgᵢ[i] + return nothing + end +end + +@inline fill_density!(ρg::NTuple{N}, ρgᵢ::Number, I::Vararg{Int,N}) where {N} = + setindex!(last(ρg), ρgᵢ, I...) + +@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::NTuple{3}, ::Val{2}) = + compute_density(rheology, args) .* (gᵢ[1], gᵢ[3]) +@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::NTuple{3}, ::Val{3}) = + compute_density(rheology, args) .* gᵢ +@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::Number, ::Any) = + compute_density(rheology, args) * gᵢ + +@inline compute_buoyancies( + rheology::MaterialParams, phase_ratios, args, gᵢ::NTuple{3}, ::Val{2} +) = compute_density_ratio(phase_ratios, rheology, args) .* (gᵢ[1], gᵢ[3]) +@inline compute_buoyancies( + rheology::MaterialParams, phase_ratios, args, gᵢ::NTuple{3}, ::Val{3} +) = compute_density_ratio(phase_ratios, rheology, args) .* gᵢ +@inline compute_buoyancies( + rheology::MaterialParams, phase_ratios, args, gᵢ::Number, ::Any +) = compute_density_ratio(phase_ratios, rheology, args) * gᵢ + +@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::NTuple{3}, ::Val{2}) = + fn_ratio(compute_density, rheology, phase_ratios, args) .* (gᵢ[1], gᵢ[3]) +@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::NTuple{3}, ::Val{3}) = + fn_ratio(compute_density, rheology, phase_ratios, args) .* gᵢ +@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::Number, ::Any) = + fn_ratio(compute_density, rheology, phase_ratios, args) * gᵢ """ compute_buoyancy(rheology::MaterialParams, args) @@ -91,14 +150,15 @@ Compute the buoyancy forces based on the given rheology, arguments, and phase ra end # without phase ratios -@inline update_ρg!(ρg::AbstractArray, rheology, args) = +@inline update_ρg!(ρg::Union{NTuple,AbstractArray}, rheology, args) = update_ρg!(isconstant(rheology), ρg, rheology, args) @inline update_ρg!(::ConstantDensityTrait, ρg, rheology, args) = nothing @inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = compute_ρg!(ρg, rheology, args) # with phase ratios -@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = - update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args) +@inline update_ρg!( + ρg::Union{NTuple,AbstractArray}, phase_ratios::JustPIC.PhaseRatios, rheology, args +) = update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args) @inline update_ρg!( ::ConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args ) = nothing diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 8faf8c31..9153fa29 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -527,7 +527,7 @@ function _solve!( Vx_on_Vy = @zeros(size(stokes.V.Vy)) # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_ρg!(ρg, phase_ratios, rheology, args) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) displacement2velocity!(stokes, dt, flow_bcs) @@ -557,7 +557,7 @@ function _solve!( args, ) - update_ρg!(ρg[2], phase_ratios, rheology, args) + update_ρg!(ρg, phase_ratios, rheology, args) @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 3f603bbe..75693f8b 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -208,7 +208,7 @@ function _solve!( θ = @zeros(ni...) # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_ρg!(ρg, phase_ratios, rheology, args) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) # convert displacement to velocity @@ -241,7 +241,7 @@ function _solve!( ) # Update buoyancy - update_ρg!(ρg[3], rheology, args) + update_ρg!(ρg, rheology, args) update_viscosity!( stokes, @@ -422,7 +422,7 @@ function _solve!( ητ = deepcopy(η) # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_ρg!(ρg, phase_ratios, rheology, args) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) # convert displacement to velocity @@ -457,7 +457,7 @@ function _solve!( ) # Update buoyancy - update_ρg!(ρg[end], phase_ratios, rheology, args) + update_ρg!(ρg, phase_ratios, rheology, args) # Update viscosity update_viscosity!( diff --git a/test/Project.toml b/test/Project.toml index d3c08479..019ed2ec 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,7 +21,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" AMDGPU = "0.8, 0.9, 1" CUDA = "5" CellArrays = "0.2" -GeoParams = "0.5, 0.6" +GeoParams = "0.6.7" JustPIC = "0.5.3" MPI = "0.20" ParallelStencil = "0.11, 0.12, 0.13"