Optimize IB Index Lookup

src/common/m_helper.fpp

@@ -20,7 +20,7 @@ module m_helper
         & s_int_to_str, s_transform_vec, s_transform_triangle, s_transform_model, s_swap, f_cross, f_create_transform_matrix, &
         & f_create_bbox, s_print_2D_array, f_xor, f_logical_to_int, associated_legendre, real_ylm, double_factorial, factorial, &
         & f_cut_on, f_cut_off, s_downsample_data, s_upsample_data, s_cross_product, f_unit_vector, s_prng, modmul, &
-        & f_local_rank_owns_location
+        & f_local_rank_owns_location, s_sort_int_key_value
 
 contains
 
@@ -366,6 +366,29 @@ contains
 
     end subroutine s_swap
 
+    !> Sort the key-value pair by the key
+    pure subroutine s_sort_int_key_value(keys, vals, n)
+
+        integer, dimension(:), intent(inout) :: keys, vals
+        integer, intent(in)                  :: n
+        integer                              :: i, j, key, val
+
+        do i = 2, n
+            key = keys(i); val = vals(i)
+
+            j = i
+            do while (j > 1)
+                if (keys(j - 1) <= key) exit
+                j = j - 1
+            end do
+
+            keys(j + 1:i) = keys(j:i - 1)
+            vals(j + 1:i) = vals(j:i - 1)
+            keys(j) = key; vals(j) = val
+        end do
+
+    end subroutine s_sort_int_key_value
+
     !> Create a transformation matrix.
     function f_create_transform_matrix(param, center) result(out_matrix)
 

src/simulation/m_collisions.fpp

@@ -21,17 +21,14 @@ module m_collisions
     implicit none
 
     private; public :: s_apply_collision_forces, s_initialize_collisions_module, s_finalize_collisions_module, &
-        & f_neighborhood_ranks_own_location, ib_gbl_idx_lookup, collisions_active
+        & f_neighborhood_ranks_own_location, collisions_active
     ! overlap distances for computing collisions
     integer, allocatable, dimension(:,:)  :: collision_lookup
     real(wp), allocatable, dimension(:,:) :: wall_overlap_distances
     real(wp)                              :: spring_stiffness, damping_parameter
     $:GPU_DECLARE(create='[spring_stiffness, damping_parameter]')
     $:GPU_DECLARE(create='[collision_lookup, wall_overlap_distances]')
 
-    integer, dimension(:), allocatable :: ib_gbl_idx_lookup
-    $:GPU_DECLARE(create='[ib_gbl_idx_lookup]')
-
     !> true when any IB-IB or IB-wall contact was detected on this rank since the last adaptive-dt computation
     logical :: collisions_active
 
@@ -249,7 +246,7 @@ contains
         integer, intent(out)                              :: num_considered_collisions
         integer                                           :: i, j, k, z_bound, ii, jj, kk
         integer, dimension(2)                             :: decoded_pairs
-        integer                                           :: gp_idx, gp_patch_id, neighbor_patch_id
+        integer                                           :: gp_idx, gp_patch_id, neighbor_patch_id, local_idx
         integer                                           :: pair_idx, out_idx
         logical                                           :: already_found
 
@@ -302,8 +299,10 @@ contains
             ! get the decoded pairs for checking if they exist, using ii,jj,kk as dummy indices
             call s_decode_patch_periodicity(raw_pairs(pair_idx, 1), decoded_pairs(1), ii, jj, kk)
             call s_decode_patch_periodicity(raw_pairs(pair_idx, 2), decoded_pairs(2), ii, jj, kk)
-            decoded_pairs(1) = ib_gbl_idx_lookup(decoded_pairs(1))
-            decoded_pairs(2) = ib_gbl_idx_lookup(decoded_pairs(2))
+            call s_get_neighborhood_idx(decoded_pairs(1), local_idx)
+            decoded_pairs(1) = local_idx
+            call s_get_neighborhood_idx(decoded_pairs(2), local_idx)
+            decoded_pairs(2) = local_idx
 
             ! skip self-collisions (an IB cannot collide with its own periodic image)
             if (decoded_pairs(1) == decoded_pairs(2)) cycle

src/simulation/m_ib_patches.fpp

@@ -19,11 +19,20 @@ module m_ib_patches
     use m_helper_basic
     use m_helper
     use m_mpi_common
+    use m_constants
 
     implicit none
 
     private; public :: s_apply_ib_patches, s_update_ib_rotation_matrix, s_instantiate_STL_models, s_decode_patch_periodicity, &
-        & s_encode_patch_periodicity, s_initialize_ib_airfoils, s_get_periodicities, s_get_ib_bound
+        & s_encode_patch_periodicity, s_initialize_ib_airfoils, s_get_periodicities, s_get_ib_bound, s_get_neighborhood_idx, &
+        & s_update_ib_lookup, s_compact_ib_lookup, s_merge_ib_lookup
+
+    !> lookup arrays for converting global IB indices to local indices
+    integer, dimension(num_ib_patches_max_namelist) :: ib_lookup_keys, ib_lookup_vals
+    $:GPU_DECLARE(create='[ib_lookup_keys, ib_lookup_vals]')
+
+    !> Holds each step's arrivals while they are sorted and merged in. Host only.
+    integer, dimension(num_ib_patches_max_namelist) :: ib_new_keys, ib_new_vals
 
 contains
 
@@ -687,6 +696,113 @@ contains
 
     end subroutine s_decode_patch_periodicity
 
+    !> binary search to retrieve the local IB patch index using the global index
+    subroutine s_get_neighborhood_idx(gbl_idx, neighborhood_idx, num_entries)
+
+        $:GPU_ROUTINE(parallelism='[seq]')
+
+        integer, intent(in)           :: gbl_idx
+        integer, intent(out)          :: neighborhood_idx
+        integer, intent(in), optional :: num_entries
+        integer                       :: lo, hi, mid
+
+        neighborhood_idx = -1
+        lo = 1
+        hi = num_ibs
+        if (present(num_entries)) hi = num_entries
+
+        do while (lo <= hi)
+            mid = lo + (hi - lo)/2
+            if (ib_lookup_keys(mid) == gbl_idx) then
+                neighborhood_idx = ib_lookup_vals(mid)
+                return
+            else if (ib_lookup_keys(mid) < gbl_idx) then
+                lo = mid + 1
+            else
+                hi = mid - 1
+            end if
+        end do
+
+    end subroutine s_get_neighborhood_idx
+
+    !> Completely rebuilds the ib lookup map, used at startup
+    subroutine s_update_ib_lookup()
+
+        integer :: i
+
+        @:PROHIBIT(num_ibs > num_ib_patches_max_namelist, &
+                   & "num_ibs exceeds the IB lookup capacity. Increase num_ib_patches_max_namelist.")
+
+        do i = 1, num_ibs
+            ib_lookup_keys(i) = patch_ib(i)%gbl_patch_id
+            ib_lookup_vals(i) = i
+        end do
+        call s_sort_int_key_value(ib_lookup_keys, ib_lookup_vals, num_ibs)
+
+        $:GPU_UPDATE(device='[ib_lookup_keys(1:num_ibs), ib_lookup_vals(1:num_ibs)]')
+
+    end subroutine s_update_ib_lookup
+
+    !> Drop the entries whose patches left the neighborhood and renumber the survivors onto the patch_ib slots they were compacted
+    !! into. Keys are never reordered, so the map stays sorted for free and only the values move: O(num_ibs_old) against re-sorting
+    !! the whole map.
+    subroutine s_compact_ib_lookup(old_to_new, num_ibs_old)
+
+        integer, dimension(:), intent(in) :: old_to_new  !< old patch_ib slot -> new slot, -1 if dropped
+        integer, intent(in)               :: num_ibs_old
+        integer                           :: i, k
+
+        k = 0
+        do i = 1, num_ibs_old
+            if (old_to_new(ib_lookup_vals(i)) < 0) cycle
+            k = k + 1
+            ib_lookup_keys(k) = ib_lookup_keys(i)
+            ib_lookup_vals(k) = old_to_new(ib_lookup_vals(i))
+        end do
+        @:ASSERT(k == num_ibs, 'IB lookup and patch_ib disagree on the surviving patch count')
+
+        $:GPU_UPDATE(device='[ib_lookup_keys(1:num_ibs), ib_lookup_vals(1:num_ibs)]')
+
+    end subroutine s_compact_ib_lookup
+
+    !> Fold patch_ib(num_ibs_pre+1:num_ibs) into the map: sort just the arrivals, then merge the two sorted runs downward from the
+    !! top. O(num_ibs) plus the sort of the few arrivals. They are copied out first because the runs share this array and merging in
+    !! place would overwrite entries still to be read.
+    subroutine s_merge_ib_lookup(num_ibs_pre)
+
+        integer, intent(in) :: num_ibs_pre
+        integer             :: i, j, k, r
+        logical             :: take_old
+
+        r = num_ibs - num_ibs_pre
+        if (r <= 0) return
+
+        do i = 1, r
+            ib_new_keys(i) = patch_ib(num_ibs_pre + i)%gbl_patch_id
+            ib_new_vals(i) = num_ibs_pre + i
+        end do
+        call s_sort_int_key_value(ib_new_keys, ib_new_vals, r)
+
+        i = num_ibs_pre; j = r; k = num_ibs
+        do while (j >= 1)
+            ! Fortran does not short-circuit .and., so the exhausted-head test stands on its own
+            take_old = .false.
+            if (i >= 1) take_old = ib_lookup_keys(i) > ib_new_keys(j)
+
+            if (take_old) then
+                ib_lookup_keys(k) = ib_lookup_keys(i); ib_lookup_vals(k) = ib_lookup_vals(i)
+                i = i - 1
+            else
+                ib_lookup_keys(k) = ib_new_keys(j); ib_lookup_vals(k) = ib_new_vals(j)
+                j = j - 1
+            end if
+            k = k - 1
+        end do
+
+        $:GPU_UPDATE(device='[ib_lookup_keys(1:num_ibs), ib_lookup_vals(1:num_ibs)]')
+
+    end subroutine s_merge_ib_lookup
+
     !> Determine the periodic wrapping bounds in each direction
     subroutine s_get_periodicities(xp_lower, xp_upper, yp_lower, yp_upper, zp_lower, zp_upper)
 

src/simulation/m_ibm.fpp

@@ -1519,16 +1519,17 @@ contains
 
     subroutine s_handoff_ib_ownership()
 
-        integer                               :: i, j, k, output_idx, local_output_idx
-        integer                               :: old_num_local_ibs
-        integer                               :: new_count, recv_count
-        integer                               :: pack_pos, unpack_pos, buf_size, patch_bytes
-        integer                               :: send_neighbor, recv_neighbor, ierr
-        integer                               :: dx, dy, dz, tag, nbr_idx, nreqs
-        real(wp), dimension(3)                :: centroid
-        logical                               :: is_new
-        type(ib_patch_parameters)             :: tmp_patch
-        integer, dimension(num_local_ibs_max) :: local_ib_idx_old
+        integer                                         :: i, j, k, output_idx, local_output_idx
+        integer                                         :: old_num_local_ibs, num_ibs_old, num_ibs_pre
+        integer                                         :: new_count, recv_count
+        integer                                         :: pack_pos, unpack_pos, buf_size, patch_bytes
+        integer                                         :: send_neighbor, recv_neighbor, ierr
+        integer                                         :: dx, dy, dz, tag, nbr_idx, nreqs
+        real(wp), dimension(3)                          :: centroid
+        logical                                         :: is_new
+        type(ib_patch_parameters)                       :: tmp_patch
+        integer, dimension(num_local_ibs_max)           :: local_ib_idx_old
+        integer, dimension(num_ib_patches_max_namelist) :: old_to_new  ! old patch_ib slot -> slot after compaction
         ! 26 neighbors max in 3D (8 in 2D); each gets its own recv buffer
         integer, parameter             :: max_nbrs = 26
         character(len=1), allocatable  :: send_buf(:), recv_bufs(:,:)
@@ -1549,15 +1550,18 @@ contains
             $:GPU_UPDATE(host='[patch_ib]')
 
             ! delete any particles that no longer need to be tracked and coalesce the array
+            num_ibs_old = num_ibs
             output_idx = 0
             local_output_idx = 0
             do i = 1, num_ibs
+                old_to_new(i) = -1
                 centroid = [patch_ib(i)%x_centroid, patch_ib(i)%y_centroid, 0._wp]
                 if (num_dims == 3) centroid(3) = patch_ib(i)%z_centroid
 
                 ! delete if not in neighborhood
                 if (f_neighborhood_ranks_own_location(centroid)) then
                     output_idx = output_idx + 1
+                    old_to_new(i) = output_idx
                     if (i /= output_idx) then
                         patch_ib(output_idx) = patch_ib(i)
                     end if
@@ -1572,7 +1576,7 @@ contains
             num_ibs = output_idx
             num_local_ibs = local_output_idx
             $:GPU_UPDATE(device='[patch_ib]')
-            call s_update_ib_lookup()
+            call s_compact_ib_lookup(old_to_new, num_ibs_old)
 
             ! Broadcast newly-owned patches to all neighborhood neighbors
             patch_bytes = storage_size(tmp_patch)/8
@@ -1641,14 +1645,15 @@ contains
             call MPI_WAITALL(nreqs, requests, MPI_STATUSES_IGNORE, ierr)
 
             ! Unpack all received buffers
+            num_ibs_pre = num_ibs
             do nbr_idx = 1, merge(26, 8, num_dims == 3)
                 if (recv_neighbor_list(nbr_idx) == MPI_PROC_NULL) cycle
                 unpack_pos = 0
                 call MPI_UNPACK(recv_bufs(:,nbr_idx), buf_size, unpack_pos, recv_count, 1, MPI_INTEGER, MPI_COMM_WORLD, ierr)
                 do i = 1, recv_count
                     call MPI_UNPACK(recv_bufs(:,nbr_idx), buf_size, unpack_pos, tmp_patch, patch_bytes, MPI_BYTE, MPI_COMM_WORLD, &
                                     & ierr)
-                    call s_get_neighborhood_idx(tmp_patch%gbl_patch_id, j)
+                    call s_get_neighborhood_idx(tmp_patch%gbl_patch_id, j, num_ibs_pre)
                     if (j < 0) then
                         num_ibs = num_ibs + 1
                         @:ASSERT(num_ibs <= size(patch_ib), 'patch_ib overflow in neighborhood handoff')
@@ -1659,48 +1664,18 @@ contains
 
             deallocate (send_buf, recv_bufs)
             $:GPU_UPDATE(device='[patch_ib]')
-            call s_update_ib_lookup()
+            call s_merge_ib_lookup(num_ibs_pre)
         end if
 #endif
 
     end subroutine s_handoff_ib_ownership
 
-    subroutine s_get_neighborhood_idx(gbl_idx, neighborhood_idx)
-
-        $:GPU_ROUTINE(parallelism='[seq]')
-
-        integer, intent(in)  :: gbl_idx
-        integer, intent(out) :: neighborhood_idx
-        integer              :: i
-
-        neighborhood_idx = ib_gbl_idx_lookup(gbl_idx)
-
-    end subroutine s_get_neighborhood_idx
-
-    subroutine s_update_ib_lookup()
-
-        integer :: i
-
-        ib_gbl_idx_lookup = -1
-        $:GPU_UPDATE(device='[ib_gbl_idx_lookup]')
-
-        $:GPU_PARALLEL_LOOP(private='[i]')
-        do i = 1, num_ibs
-            ib_gbl_idx_lookup(patch_ib(i)%gbl_patch_id) = i
-        end do
-        $:END_GPU_PARALLEL_LOOP()
-
-        $:GPU_UPDATE(host='[ib_gbl_idx_lookup]')
-
-    end subroutine s_update_ib_lookup
-
     !> Finalize the IBM module
     impure subroutine s_finalize_ibm_module()
 
         integer :: i
 
         @:DEALLOCATE(ib_markers%sf)
-        @:DEALLOCATE(ib_gbl_idx_lookup)
         do i = 1, num_ib_airfoils_max
             if (allocated(ib_airfoil_grids(i)%upper)) then
                 @:DEALLOCATE(ib_airfoil_grids(i)%upper)

src/simulation/m_start_up.fpp

@@ -1337,8 +1337,6 @@ contains
         end if
 #endif
 
-        @:ALLOCATE(ib_gbl_idx_lookup(1:num_gbl_ibs))
-
     end subroutine s_build_ib_neighborhood
 
     !> Build ib_neighbor_ranks(-1:1,-1:1,-1:1): MPI ranks of all neighbor domains. Uses two rounds of MPI_SENDRECV cascades - face