Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adapts WENO5 struct for GPU #251

Merged
merged 19 commits into from
Oct 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,12 @@ on:
pull_request:
branches:
- main

# Cancel redundant CI tests automatically
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true

jobs:
test:
name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
# ----------------------------------------------------

# Weno model -----------------------------------------
weno = WENO5(Val(2), (nx,ny).+1) # ni.+1 for Temp
weno = WENO5(backend, Val(2), (nx,ny).+1) # ni.+1 for Temp
# ----------------------------------------------------

# Initialize particles -------------------------------
Expand Down
2 changes: 1 addition & 1 deletion miniapps/convection/GlobalConvection2D_WENO5.jl
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma
# ----------------------------------------------------

# Weno model -----------------------------------------
weno = WENO5(Val(2), ni.+1) # ni.+1 for Temp
weno = WENO5(backend, Val(2), ni.+1) # ni.+1 for Temp
# ----------------------------------------------------

# create rheology struct
Expand Down
6 changes: 3 additions & 3 deletions miniapps/convection/WENO5/WENO_convection2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
# Weno model -----------------------------------------
weno = WENO5(backend_JR, Val(2), ni.+1)
# ----------------------------------------------------

# Initialize particles -------------------------------
nxcell, max_xcell, min_xcell = 20, 40, 1
particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi...);
Expand Down Expand Up @@ -285,8 +285,8 @@ end
figdir = "Weno2D"
ar = 1 # aspect ratio
n = 64
nx = n*ar
ny = n
nx = n*ar
ny = n
igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid
IGG(init_global_grid(nx, ny, 1; init_MPI= true)...)
else
Expand Down
4 changes: 4 additions & 0 deletions src/ext/AMDGPU/2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -364,4 +364,8 @@ function JR2D.compute_shear_heating!(
return nothing
end

function JR2D.WENO_advection!(u::ROCArrayArray, Vxi::NTuple, weno, di, dt)
return WENO_advection!(u, Vxi, weno, di, dt)
end

end
9 changes: 5 additions & 4 deletions src/ext/AMDGPU/3D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -236,10 +236,6 @@ function JR3D.compute_melt_fraction!(
return compute_melt_fraction!(ϕ, phase_ratios, rheology, args)
end

function JR3D.compute_melt_fraction!(ϕ::ROCArray, phase_ratios, rheology, args)
return compute_melt_fraction!(ϕ, phase_ratios, rheology, args)
end

# Interpolations
function JR3D.temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays)
return _temperature2center!(thermal)
Expand Down Expand Up @@ -385,4 +381,9 @@ function JR3D.compute_shear_heating!(
)
return nothing
end

function JR3D.WENO_advection!(u::ROCArray, Vxi::NTuple, weno, di, dt)
return WENO_advection!(u, Vxi, weno, di, dt)
end

end
5 changes: 5 additions & 0 deletions src/ext/CUDA/2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -348,4 +348,9 @@ function JR2D.compute_shear_heating!(
return nothing
end

# WENO-5 advection
function JR2D.WENO_advection!(u::CuArray, Vxi::NTuple, weno, di, dt)
return WENO_advection!(u, Vxi, weno, di, dt)
end

end
22 changes: 4 additions & 18 deletions src/ext/CUDA/3D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -180,20 +180,6 @@ end
function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs)
return thermal_bcs!(thermal.T, bcs)
end

# # Phases
# function JR3D.phase_ratios_center!(
# ::CUDABackendTrait, phase_ratios::JustPIC.PhaseRatios, particles, grid::Geometry, phases
# )
# return _phase_ratios_center!(phase_ratios, particles, grid, phases)
# end

# function JR3D.phase_ratios_vertex!(
# ::CUDABackendTrait, phase_ratios::JustPIC.PhaseRatios, particles, grid::Geometry, phases
# )
# return _phase_ratios_vertex!(phase_ratios, particles, grid, phases)
# end

# Rheology

## viscosity
Expand Down Expand Up @@ -252,10 +238,6 @@ function JR3D.compute_melt_fraction!(
return compute_melt_fraction!(ϕ, phase_ratios, rheology, args)
end

function JR3D.compute_melt_fraction!(ϕ::CuArray, phase_ratios, rheology, args)
return compute_melt_fraction!(ϕ, phase_ratios, rheology, args)
end

# Interpolations
function JR3D.temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays)
return _temperature2center!(thermal)
Expand Down Expand Up @@ -399,4 +381,8 @@ function JR3D.compute_shear_heating!(
return nothing
end

function JR3D.WENO_advection!(u::CuArray, Vxi::NTuple, weno, di, dt)
return WENO_advection!(u, Vxi, weno, di, dt)
end

end
19 changes: 7 additions & 12 deletions src/rheology/Melting.jl
Original file line number Diff line number Diff line change
@@ -1,31 +1,26 @@
function compute_melt_fraction!(ϕ, rheology, args)
ni = size(ϕ)
@parallel (@idx ni) compute_melt_fraction_kernel!(ϕ, rheology, args)
return nothing
end

@parallel_indices (I...) function compute_melt_fraction_kernel!(ϕ, rheology, args)
args_ijk = ntuple_idx(args, I...)
ϕ[I...] = compute_melt_frac(rheology, args_ijk)
ϕ[I...] = compute_meltfraction(rheology, args_ijk)
return nothing
end

@inline function compute_melt_frac(rheology, args)
return compute_meltfraction(rheology, args)
end

function compute_melt_fraction!(ϕ, phase_ratios, rheology, args)
function compute_melt_fraction!(ϕ, phase_ratios::JustPIC.PhaseRatios, rheology, args)
ni = size(ϕ)
@parallel (@idx ni) compute_melt_fraction_kernel!(ϕ, phase_ratios, rheology, args)
@parallel (@idx ni) compute_melt_fraction_kernel!(
ϕ, phase_ratios.center, rheology, args
)
end

@parallel_indices (I...) function compute_melt_fraction_kernel!(
ϕ, phase_ratios, rheology, args
)
args_ijk = ntuple_idx(args, I...)
ϕ[I...] = compute_melt_frac(rheology, args_ijk, phase_ratios[I...])
ϕ[I...] = fn_ratio(compute_meltfraction, rheology, @cell(phase_ratios[I...]), args_ijk)
return nothing
end

@inline function compute_melt_frac(rheology, args, phase_ratios)
return compute_meltfraction_ratio(phase_ratios, rheology, args)
end
3 changes: 3 additions & 0 deletions src/types/weno.jl
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
using Adapt
## Weno5 advection scheme. Implementation based on the repository from
# https://gmd.copernicus.org/preprints/gmd-2023-189/

Expand Down Expand Up @@ -59,3 +60,5 @@ struct WENO5{T,N,A,M} <: AbstractWENO
# method
method::M # 1:JS, 2:Z
end

Adapt.@adapt_structure WENO5
10 changes: 1 addition & 9 deletions test/test_Blankenbach.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ push!(LOAD_PATH, "..")

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using Test, Suppressor
Expand Down Expand Up @@ -137,12 +135,6 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10)
temperature2center!(thermal)
# ----------------------------------------------------

# Rayleigh number
ΔT = thermal.T[1,1] - thermal.T[1,end]
Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) /
(κ * rheology[1].CompositeRheology[1].elements[1].η )
@show Ra

args = (; T = thermal.Tc, P = stokes.P, dt = Inf)

# Buoyancy forces & viscosity ----------------------
Expand Down Expand Up @@ -264,7 +256,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10)
# Compute U rms -----------------------------
# U₍ᵣₘₛ₎ = H*ρ₀*c₍ₚ₎/k * √ 1/H/L * ∫∫ (vx²+vz²) dx dz
Urms_it = let
JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true)
velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true)
@. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude
sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) *
((ly * rheology[1].Density[1].ρ0 * rheology[1].HeatCapacity[1].Cp) / rheology[1].Conductivity[1].k )
Expand Down
7 changes: 4 additions & 3 deletions test/test_Interpolations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ push!(LOAD_PATH, "..")

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using Test
Expand All @@ -27,7 +25,7 @@ end


@testset "Interpolations" begin

if backend_JR == CPUBackend
# Set up mock data
# Physical domain ------------------------------------
ly = 1.0 # domain length in y
Expand Down Expand Up @@ -75,4 +73,7 @@ end
velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=false)
@test iszero(Vx_v[1,1])
@test Vy_v[1,1] == 10
else
@test true == true
end
end
4 changes: 1 addition & 3 deletions test/test_WENO5.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ push!(LOAD_PATH, "..")

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using Test, Suppressor
Expand Down Expand Up @@ -112,7 +110,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, thermal_perturbation =
# ----------------------------------------------------

# Weno model -----------------------------------------
weno = WENO5(Val(2), ni.+1) # ni.+1 for Temp
weno = WENO5(backend_JR, Val(2), ni.+1) # ni.+1 for Temp
# ----------------------------------------------------

# create rheology struct
Expand Down
2 changes: 0 additions & 2 deletions test/test_boundary_conditions2D.jl
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using JustRelax, JustRelax.JustRelax2D
Expand Down
2 changes: 0 additions & 2 deletions test/test_boundary_conditions3D.jl
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using JustRelax, JustRelax.JustRelax3D
Expand Down
2 changes: 0 additions & 2 deletions test/test_shearheating2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,8 @@ using Test, Suppressor

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

# Benchmark of Duretz et al. 2014
Expand Down
2 changes: 0 additions & 2 deletions test/test_shearheating3D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,8 @@ using Test, Suppressor

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

# Benchmark of Duretz et al. 2014
Expand Down
2 changes: 0 additions & 2 deletions test/test_sinking_block.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ push!(LOAD_PATH, "..")

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true) #used because of velocity2vertex!
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true) #used because of velocity2vertex!
end


Expand Down
2 changes: 0 additions & 2 deletions test/test_stokes_solvi3D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@ push!(LOAD_PATH, "..")

@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
using AMDGPU
AMDGPU.allowscalar(true)
elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
using CUDA
CUDA.allowscalar(true)
end

using Test, Suppressor
Expand Down
Loading