From d3b9e85d64c2e111f818fac75c8120e15f6d9a23 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Fri, 16 Aug 2024 14:28:57 +0200 Subject: [PATCH] update PT coefs inside iterations --- src/ext/AMDGPU/2D.jl | 24 +++++++++++++++++++ src/ext/AMDGPU/3D.jl | 24 +++++++++++++++++++ src/ext/CUDA/2D.jl | 24 +++++++++++++++++++ src/ext/CUDA/3D.jl | 24 +++++++++++++++++++ .../DiffusionPT_coefficients.jl | 24 +++++++++++++++++++ src/thermal_diffusion/DiffusionPT_solver.jl | 1 + 6 files changed, 121 insertions(+) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 28093b54..aef41d40 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -85,6 +85,30 @@ function JR2D.PTThermalCoeffs( return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) end +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, phase_ratios, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, ::Nothing, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + # Boundary conditions function JR2D.flow_bcs!( ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 112347d2..c03ba3b7 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -85,6 +85,30 @@ function JR3D.PTThermalCoeffs( return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) end +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, phase_ratios, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:RocArray}, rheology, ::Nothing, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + # Boundary conditions function JR3D.flow_bcs!( ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index cbe74022..981f3e9d 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -84,6 +84,30 @@ function JR2D.PTThermalCoeffs( return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) end +function JR2D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR2D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR2D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + # Boundary conditions function JR2D.flow_bcs!( ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 7013245a..7056f533 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -95,6 +95,30 @@ function JR3D.update_pt_thermal_arrays!( return nothing end +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D. update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function JR3D.update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt) where {T} + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + # Boundary conditions function JR3D.flow_bcs!( ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index acd809c0..793f8fa2 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -121,3 +121,27 @@ function _compute_pt_thermal_arrays!( return nothing end + +function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, phase_ratios, args, dt) + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, phase_ratios.center, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, args, dt) + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end + +function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, ::Nothing, args, dt) + ni = size(pt_thermal.dτ_ρ) + @parallel (@idx ni) JustRelax2D.compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, pt_thermal.dτ_ρ, rheology, args, pt_thermal.max_lxyz, pt_thermal.Vpdτ, inv(dt) + ) + return nothing +end \ No newline at end of file diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index dc4b76df..5f2bb886 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -147,6 +147,7 @@ function _heatdiffusion_PT!( while err > ϵ && iter < iterMax wtime0 += @elapsed begin + update_thermal_coeffs!(pt_thermal, rheology, phase, args, dt) @parallel flux_range(ni...) compute_flux!( @qT(thermal)..., @qT2(thermal)...,