Merge pull request #171 from JuliaGeodynamics/adm/move

Improve `move!` kernel performance

src/Particles/move_safe.jl

@@ -26,8 +26,11 @@ function move_particles!(particles::AbstractParticles, grid::NTuple{N}, args) wh
             )
         end
     elseif N == 3
+        nthreads = (16, 16, 1)
+        # nthreads = (6, 6, 6)
+        nblocks  = ceil.(Int, n_color ./ nthreads)
         for offsetᵢ in 1:3, offsetⱼ in 1:3, offsetₖ in 1:3
-            @parallel (@idx n_color) move_particles_ps!(
+            @parallel (@idx n_color) nblocks nthreads move_particles_ps!(
                 coords, grid, dxi, index, domain_limits, args, (offsetᵢ, offsetⱼ, offsetₖ)
             )
         end
@@ -70,6 +73,7 @@ function move_kernel!(
     idx::NTuple{N1,Int64},
 ) where {N1,N2,T}
 
+    starting_point = 1
     # iterate over particles in child cell 
     for ip in cellaxes(index)
         doskip(index, ip, idx...) && continue
@@ -102,8 +106,10 @@ function move_kernel!(
         empty_particle!(coords, ip, idx)
         empty_particle!(args, ip, idx)
         # check whether there's empty space in parent cell
-        free_idx = find_free_memory(index, new_cell...)
+        # free_idx = find_free_memory(index, new_cell...)
+        free_idx = find_free_memory(starting_point, index, new_cell...)
         free_idx == 0 && continue
+        starting_point = free_idx
         # move particle and its fields to the first free memory location
         @inbounds @index index[free_idx, new_cell...] = true
         fill_particle!(coords, pᵢ, free_idx, new_cell)
@@ -131,7 +137,14 @@ end
 
 function find_free_memory(index, I::Vararg{Int,N}) where {N}
     for i in cellaxes(index)
-        !(@index(index[i, I...])) && return i
+        (@index(index[i, I...])) || return i
+    end
+    return 0
+end
+
+function find_free_memory(initial_index::Integer, index, I::Vararg{Int,N}) where {N}
+    for i in initial_index:cellnum(index)
+        (@index(index[i, I...])) || return i
     end
     return 0
 end
@@ -140,7 +153,7 @@ end
     quote
         Base.@_inline_meta
         Base.Cartesian.@nexprs $N i ->
-            (@inbounds !(domain_limits[i][1] < p[i] < domain_limits[i][2]) && return false)
+            (@inbounds (domain_limits[i][1] < p[i] < domain_limits[i][2]) || return false)
         return true
     end
 end

test/test_2D.jl

@@ -266,8 +266,8 @@ end
 
 function advection_test_2D()
     # Initialize particles -------------------------------
-    nxcell, max_xcell, min_xcell = 25, 40, 10
-    n = 64
+    nxcell, max_xcell, min_xcell = 25, 50, 10
+    n = 128
     nx = ny = n-1
     Lx = Ly = 1.0
     # nodal vertices
@@ -324,8 +324,8 @@ end
 
 function test_rotating_circle()
     # Initialize particles -------------------------------
-    nxcell, max_xcell, min_xcell = 50, 60, 40
-    n = 101
+    nxcell, max_xcell, min_xcell = 25, 50, 10
+    n = 256
     nx = ny = n-1
     Lx = Ly = 1.0
     # nodal vertices
@@ -364,7 +364,7 @@ function test_rotating_circle()
     while t ≤ tmax
         _2D.particle2grid!(T, pT, xvi, particles)
         copyto!(T0, T)
-        _2D.advection!(particles, _2D.RungeKutta2(2/3), V, (grid_vx, grid_vy), dt)
+        _2D.advection!(particles, _2D.RungeKutta2(), V, (grid_vx, grid_vy), dt)
         _2D.move_particles!(particles, xvi, particle_args)
         _2D.inject_particles!(particles, (pT, ), xvi)
         _2D.grid2particle!(pT, xvi, T, particles)

test/test_3D.jl

@@ -201,7 +201,7 @@ end
 
 function test_advection_3D()
 
-    n   = 80
+    n   = 64
     nx  = ny = nz = n-1
     Lx  = Ly = Lz = 1.0
     ni  = nx, ny, nz
@@ -265,9 +265,6 @@ function test_advection()
     return passed
 end
 
-env = ENV["JULIA_JUSTPIC_BACKEND"]
-if !(env === "AMDGPU" && env === "CUDA")
-    @testset "Miniapps" begin
-        @test test_advection()
-    end
+@testset "Miniapps" begin
+    @test test_advection()
 end