From f50bab36a5579f264b8367054f5d92fb5a1dee15 Mon Sep 17 00:00:00 2001
From: Albert de Montserrat
<58044444+albert-de-montserrat@users.noreply.github.com>
Date: Mon, 6 May 2024 15:20:41 +0200
Subject: [PATCH] Update 3D.jl
---
src/ext/AMDGPU/3D.jl | 40 ++++++++++++++++++++++++++++++++++++++--
1 file changed, 38 insertions(+), 2 deletions(-)
diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
index 472905e1..acf39f5a 100644
--- a/src/ext/AMDGPU/3D.jl
+++ b/src/ext/AMDGPU/3D.jl
@@ -92,14 +92,14 @@ function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs)
end
# Phases
-function JR3D.phase_ratios_center(
+function JR3D.phase_ratios_center!(
::CUDABackendTrait,
phase_ratios::JustRelax.PhaseRatio,
particles,
grid::Geometry,
phases,
)
- return _phase_ratios_center(phase_ratios, particles, grid, phases)
+ return _phase_ratios_center!(phase_ratios, particles, grid, phases)
end
# Rheology
@@ -183,4 +183,40 @@ function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...)
return _compute_dt(S, args...)
end
+function JR3D.subgrid_characteristic_time!(
+ subgrid_arrays,
+ particles,
+ dt₀::RocArray,
+ phases::JustRelax.PhaseRatio,
+ rheology,
+ thermal::JustRelax.ThermalArrays,
+ stokes::JustRelax.StokesArrays,
+ xci,
+ di,
+)
+ ni = size(stokes.P)
+ @parallel (@idx ni) subgrid_characteristic_time!(
+ dt₀, phases.center, rheology, thermal.Tc, stokes.P, di
+ )
+ return nothing
+end
+
+function JR3D.subgrid_characteristic_time!(
+ subgrid_arrays,
+ particles,
+ dt₀::RocArray,
+ phases::AbstractArray{Int,N},
+ rheology,
+ thermal::JustRelax.ThermalArrays,
+ stokes::JustRelax.StokesArrays,
+ xci,
+ di,
+) where {N}
+ ni = size(stokes.P)
+ @parallel (@idx ni) subgrid_characteristic_time!(
+ dt₀, phases, rheology, thermal.Tc, stokes.P, di
+ )
+ return nothing
+end
+
end