update benchmarks

PTsolvers · Jan 30, 2024 · 51b2a62 · 51b2a62
1 parent 81f3dd4
commit 51b2a62
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 161 deletions.
diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl
@@ -1,10 +1,12 @@
 # Benchmark of Duretz et al. 2014
 # http://dx.doi.org/10.1002/2014GL060438
-using JustRelax, JustRelax.DataIO, JustPIC
+using JustRelax, JustRelax.DataIO
 import JustRelax.@cell
+
 ## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
-# set_backend("Threads_Float64_2D")
-# set_backend("CUDA_Float64_2D")
+const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+using JustPIC
+using JustPIC._2D
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:cpu, Float64, 2)
@@ -24,32 +26,6 @@ include("Shearheating_rheology.jl")
 @inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
 @inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))
 
-function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, nx, ny)
-    ni     = nx, ny
-    ncells = nx * ny
-    np     = max_xcell * ncells
-    px, py = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(2))
-    inject = @fill(false, nx, ny, eltype=Bool)
-    index  = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool)
-
-    @parallel_indices (i, j) function fill_coords_index(px, py, index)
-        # lower-left corner of the cell
-        x0, y0 = x[i], y[j]
-        # fill index array
-        for l in 1:nxcell
-            JustRelax.@cell px[l, i, j]    = x0 + dx * rand(0.05:1e-5:0.95)
-            JustRelax.@cell py[l, i, j]    = y0 + dy * rand(0.05:1e-5:0.95)
-            JustRelax.@cell index[l, i, j] = true
-        end
-        return nothing
-    end
-
-    @parallel (1:nx, 1:ny) fill_coords_index(px, py, index)
-
-    return Particles(
-        (px, py), index, inject, nxcell, max_xcell, min_xcell, np, (nx, ny)
-    )
-end
 
 # Velocity helper grids for the particle advection
 function velocity_grids(xci, xvi, di)
@@ -109,8 +85,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 20, 40, 1
-    particles = init_particles_cellarrays(
-        nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni
     )
     # velocity grids
     grid_vx, grid_vy = velocity_grids(xci, xvi, di)

diff --git a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl
@@ -1,12 +1,13 @@
 # Benchmark of Duretz et al. 2014
 # http://dx.doi.org/10.1002/2014GL060438
-using JustRelax, JustRelax.DataIO, JustPIC
-import JustRelax.@cell
 using CUDA
+using JustRelax, JustRelax.DataIO
+import JustRelax.@cell
 
 ## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
-# set_backend("Threads_Float64_3D")
-# set_backend("CUDA_Float64_3D")
+const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+using JustPIC
+using JustPIC._3D
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:cpu, Float64, 3)
@@ -25,34 +26,6 @@ include("Shearheating_rheology.jl")
 @inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
 @inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))
 
-function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, z, dx, dy, dz, ni::NTuple{3, Int})
-    ncells     = prod(ni)
-    np         = max_xcell * ncells
-    px, py, pz = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(3))
-    inject     = @fill(false, ni..., eltype=Bool)
-    index      = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool)
-
-    @parallel_indices (i, j, k) function fill_coords_index(px, py, pz, index)
-        @inline r()= rand(0.05:1e-5:0.95)
-        I          = i, j, k
-        # lower-left corner of the cell
-        x0, y0, z0 = x[i], y[j], z[k]
-        # fill index array
-        for l in 1:nxcell
-            JustRelax.@cell px[l, I...]    = x0 + dx * r()
-            JustRelax.@cell py[l, I...]    = y0 + dy * r()
-            JustRelax.@cell pz[l, I...]    = z0 + dz * r()
-            JustRelax.@cell index[l, I...] = true
-        end
-        return nothing
-    end
-
-    @parallel (@idx ni) fill_coords_index(px, py, pz, index)
-
-    return Particles(
-        (px, py, pz), index, inject, nxcell, max_xcell, min_xcell, np, ni
-    )
-end
 
 # Velocity helper grids for the particle advection
 function velocity_grids(xci, xvi, di)
@@ -102,8 +75,8 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, figdir="figs3D", save_vtk =false)
 
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 20, 40, 1
-    particles                    = init_particles_cellarrays(
-        nxcell, max_xcell, min_xcell, xvi..., di..., ni
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni
     )
     # velocity grids
     grid_vx, grid_vy, grid_vz   = velocity_grids(xci, xvi, di)

diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl
@@ -1,9 +1,10 @@
-using JustRelax, JustRelax.DataIO, JustPIC
+using JustRelax, JustRelax.DataIO
 import JustRelax.@cell
 
 ## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
-# set_backend("Threads_Float64_2D")
-# set_backend("CUDA_Float64_2D")
+const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+using JustPIC
+using JustPIC._2D
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:cpu, Float64, 2)
@@ -22,32 +23,6 @@ include("Layered_rheology.jl")
 @inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
 @inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))
 
-function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, nx, ny)
-    ni     = nx, ny
-    ncells = nx * ny
-    np     = max_xcell * ncells
-    px, py = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(2))
-    inject = @fill(false, nx, ny, eltype=Bool)
-    index  = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool)
-
-    @parallel_indices (i, j) function fill_coords_index(px, py, index)
-        # lower-left corner of the cell
-        x0, y0 = x[i], y[j]
-        # fill index array
-        for l in 1:nxcell
-            JustRelax.@cell px[l, i, j]    = x0 + dx * rand(0.05:1e-5:0.95)
-            JustRelax.@cell py[l, i, j]    = y0 + dy * rand(0.05:1e-5:0.95)
-            JustRelax.@cell index[l, i, j] = true
-        end
-        return nothing
-    end
-
-    @parallel (1:nx, 1:ny) fill_coords_index(px, py, index)
-
-    return Particles(
-        (px, py), index, inject, nxcell, max_xcell, min_xcell, np, (nx, ny)
-    )
-end
 
 # Velocity helper grids for the particle advection
 function velocity_grids(xci, xvi, di)
@@ -152,8 +127,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 20, 40, 1
-    particles = init_particles_cellarrays(
-        nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni
     )
     # velocity grids
     grid_vx, grid_vy = velocity_grids(xci, xvi, di)

diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl
@@ -1,9 +1,10 @@
-using JustRelax, JustRelax.DataIO, JustPIC
+using JustRelax, JustRelax.DataIO
 import JustRelax.@cell
 
 ## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
-# set_backend("Threads_Float64_2D")
-# set_backend("CUDA_Float64_2D")
+const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+using JustPIC
+using JustPIC._3D
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:CUDA, Float64, 3) # or (:Threads, Float64, 3) or (:AMDGPU, Float64, 3)
@@ -21,35 +22,6 @@ include("Layered_rheology.jl")
 @inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
 @inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))
 
-function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, z, dx, dy, dz, ni::NTuple{3, Int})
-    ncells     = prod(ni)
-    np         = max_xcell * ncells
-    px, py, pz = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(3))
-    inject     = @fill(false, ni..., eltype=Bool)
-    index      = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool)
-
-    @parallel_indices (i, j, k) function fill_coords_index(px, py, pz, index)
-        @inline r()= rand(0.05:1e-5:0.95)
-        I          = i, j, k
-        # lower-left corner of the cell
-        x0, y0, z0 = x[i], y[j], z[k]
-        # fill index array
-        for l in 1:nxcell
-            JustRelax.@cell px[l, I...]    = x0 + dx * r()
-            JustRelax.@cell py[l, I...]    = y0 + dy * r()
-            JustRelax.@cell pz[l, I...]    = z0 + dz * r()
-            JustRelax.@cell index[l, I...] = true
-        end
-        return nothing
-    end
-
-    @parallel (@idx ni) fill_coords_index(px, py, pz, index)
-
-    return Particles(
-        (px, py, pz), index, inject, nxcell, max_xcell, min_xcell, np, ni
-    )
-end
-
 # Velocity helper grids for the particle advection
 function velocity_grids(xci, xvi, di)
     xghost  = ntuple(Val(3)) do i
@@ -140,8 +112,8 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
 
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 20, 20, 1
-    particles                    = init_particles_cellarrays(
-        nxcell, max_xcell, min_xcell, xvi..., di..., ni
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni
     )
     # velocity grids
     grid_vx, grid_vy, grid_vz   = velocity_grids(xci, xvi, di)

diff --git a/miniapps/convection/WENO5/WENO_convection2D.jl b/miniapps/convection/WENO5/WENO_convection2D.jl
@@ -1,12 +1,13 @@
 using CUDA
 CUDA.allowscalar(false) # for safety
 
-using JustRelax, JustRelax.DataIO, JustPIC
+using JustRelax, JustRelax.DataIO
 import JustRelax.@cell
 
 ## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
-# set_backend("Threads_Float64_2D")
-# set_backend("CUDA_Float64_2D")
+const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+using JustPIC
+using JustPIC._2D
 
 # setup ParallelStencil.jl environment
 model = PS_Setup(:CUDA, Float64, 2)
@@ -24,32 +25,6 @@ include("Layered_rheology.jl")
 @inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
 @inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))
 
-function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, nx, ny)
-    ni     = nx, ny
-    ncells = nx * ny
-    np     = max_xcell * ncells
-    px, py = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(2))
-    inject = @fill(false, nx, ny, eltype=Bool)
-    index  = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool)
-
-    @parallel_indices (i, j) function fill_coords_index(px, py, index)
-        # lower-left corner of the cell
-        x0, y0 = x[i], y[j]
-        # fill index array
-        for l in 1:nxcell
-            JustRelax.@cell px[l, i, j]    = x0 + dx * rand(0.05:1e-5:0.95)
-            JustRelax.@cell py[l, i, j]    = y0 + dy * rand(0.05:1e-5:0.95)
-            JustRelax.@cell index[l, i, j] = true
-        end
-        return nothing
-    end
-
-    @parallel (1:nx, 1:ny) fill_coords_index(px, py, index)
-
-    return Particles(
-        (px, py), index, inject, nxcell, max_xcell, min_xcell, np, (nx, ny)
-    )
-end
 
 # Velocity helper grids for the particle advection
 function velocity_grids(xci, xvi, di)
@@ -146,8 +121,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 20, 40, 1
-    particles = init_particles_cellarrays(
-        nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny
+    particles = init_particles(
+        backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni
     )
     # velocity grids
     grid_vx, grid_vy = velocity_grids(xci, xvi, di)