Optimize IB Index Lookup

src/common/m_constants.fpp

@@ -35,8 +35,8 @@ module m_constants
     integer, parameter  :: num_stl_models_max = 10
     !> Maximum number of immersed boundary patches (legacy, not used for patch_ib sizing)
     !> Fixed capacity of patch_ib (namelist patches + local particle bed subset after reduction)
-    integer, parameter  :: num_local_ibs_max = 8000            !< Maximum number of immersed boundary patches (patch_ib)
-    integer, parameter  :: num_ib_patches_max_namelist = 216000
+    integer, parameter  :: num_local_ibs_max = 30000           !< Maximum number of immersed boundary patches (patch_ib)
+    integer, parameter  :: num_ib_patches_max_namelist = 810000
     integer, parameter  :: num_particle_clouds_max = 10        !< Maximum number of particle bed patch specifications
     integer, parameter  :: num_bc_patches_max = 10             !< Maximum number of boundary condition patches
     integer, parameter  :: max_2d_fourier_modes = 10           !< Max Fourier mode index for 2D modal patch (geometry 13)

src/common/m_derived_types.fpp

@@ -340,8 +340,9 @@ module m_derived_types
     end type ib_stl_parameters
 
     type ib_patch_parameters
-        integer  :: geometry                            !< Type of geometry for the patch
+        integer  :: geometry  !< Type of geometry for the patch
         integer  :: gbl_patch_id
+        integer  :: owner_rank  !< MPI rank whose subdomain holds the centroid; the only rank that sums this IB's force
         real(wp) :: x_centroid, y_centroid, z_centroid  !< Geometric center coordinates of the patch
 
         !> Centroid locations of intermediate steps in the time_stepper module

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, &
-        & s_prng_splitmix32, f_local_rank_owns_location
+        & s_prng_splitmix32, f_local_rank_owns_location, s_sort_int_key_value
 
 contains
 
@@ -395,6 +395,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
 
@@ -251,7 +248,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
 
@@ -306,8 +303,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
 
@@ -728,6 +737,117 @@ 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
+
+        call nvtxStartRange("MERGE-IB-LOOKUP")
+
+        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)]')
+
+        call nvtxEndRange()
+
+    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)