Optimize IB Index Lookup

docs/documentation/case.md

@@ -108,6 +108,7 @@ is equivalent to `"riemann_solver": 2`. Defined names appear in each parameter's
 | Parameter        | Type           | Description                               |
 | ---:             |    :----:      |          :---                             |
 | `run_time_info`  | Logical        | Output run-time information               |
+| `phase_timing_wrt` | Logical      | Append per-phase wall times to `phase_time_data.dat` |
 | `rdma_mpi`       | Logical        | (GPUs) Enable RDMA for MPI communication. |
 | `case_dir`       | String         | Case directory path                       |
 | `old_grid`       | Logical        | Use grid from previous simulation         |
@@ -116,6 +117,7 @@ is equivalent to `"riemann_solver": 2`. Defined names appear in each parameter's
 | `n_start_old`    | Integer        | Starting index from previous simulation   |
 
 - `run_time_info` generates a text file that includes run-time information including the CFL number(s) at each time-step.
+- `phase_timing_wrt` times every NVTX range of the time march on each rank and, at the end of the run, appends one row per range to `phase_time_data.dat` with the min, mean and max over ranks of its inclusive wall time per step. Nothing is communicated until the end of the run. Plot the rows with `misc/plot_phase_timings.py`.
 - `rdma_mpi` optimizes data transfers between GPUs using Remote Direct Memory Access (RDMA).
 The underlying MPI implementation and communication infrastructure must support this
 feature, detecting GPU pointers and performing RDMA accordingly.

misc/plot_phase_timings.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+"""Plot the max-over-ranks wall time per step of each phase against rank count.
+
+Reads one or more phase_time_data.dat files written by simulation with phase_timing_wrt = T
+(a directory is read as <dir>/phase_time_data.dat). Rows from all files are pooled; when
+several runs share a rank count, the last one read wins.
+"""
+
+import argparse
+import os
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+
+
+def read_rows(paths):
+    data = defaultdict(dict)  # phase -> {ranks: (mean, max)}
+    for path in paths:
+        if os.path.isdir(path):
+            path = os.path.join(path, "phase_time_data.dat")
+        with open(path) as f:
+            next(f)
+            for line in f:
+                ranks, _, _, _, t_mean, t_max, phase = line.split()
+                data[phase][int(ranks)] = (float(t_mean), float(t_max))
+    return data
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("paths", nargs="+", help="phase_time_data.dat files or the case directories holding them")
+    parser.add_argument("--phases", nargs="+", help="phases to plot (default: --top largest)")
+    parser.add_argument("--top", type=int, default=10, help="plot the N phases with the largest max time at the largest rank count")
+    parser.add_argument("--mean", action="store_true", help="also plot the mean over ranks (dashed)")
+    parser.add_argument("--relative", action="store_true", help="divide each phase by its time at the smallest rank count")
+    parser.add_argument("-o", "--output", default="phase_timings.png", help="output image")
+    args = parser.parse_args()
+
+    data = read_rows(args.paths)
+    phases = args.phases or sorted(data, key=lambda ph: data[ph][max(data[ph])][1], reverse=True)[: args.top]
+
+    fig, ax = plt.subplots(figsize=(9, 6))
+    for phase in phases:
+        ranks = sorted(data[phase])
+        t_mean, t_max = zip(*(data[phase][r] for r in ranks))
+        scale = t_max[0] if args.relative else 1.0
+        (line,) = ax.plot(ranks, [t / scale for t in t_max], "o-", label=phase)
+        if args.mean:
+            ax.plot(ranks, [t / scale for t in t_mean], "--", color=line.get_color())
+
+    ax.set_xscale("log", base=2)
+    ax.set_xlabel("Ranks")
+    ax.set_ylabel("max time / time at fewest ranks" if args.relative else "max over ranks [s/step]")
+    ax.grid(True, which="both", alpha=0.3)
+    ax.legend(fontsize="small")
+    fig.tight_layout()
+    fig.savefig(args.output, dpi=150)
+    print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

src/common/m_constants.fpp

@@ -28,8 +28,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, &
-        & 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/common/m_mpi_common.fpp

@@ -283,6 +283,63 @@ contains
 
     end subroutine mpi_bcast_time_step_values
 
+    !> Reduce the per-rank NVTX range timers onto rank 0 over the union of range names seen by any rank
+    impure subroutine s_mpi_reduce_nvtx_timers(names, num_names, t_min, t_sum, t_max, calls_max)
+
+        character(len=nvtx_name_len), intent(out)                :: names(nvtx_max_timers)
+        integer, intent(out)                                     :: num_names
+        real(kind=8), dimension(nvtx_max_timers), intent(out)    :: t_min, t_sum, t_max
+        integer(kind=8), dimension(nvtx_max_timers), intent(out) :: calls_max
+        real(kind=8)                                             :: sec(nvtx_num_timers), t_loc(nvtx_max_timers)
+        integer(kind=8)                                          :: calls_loc(nvtx_max_timers)
+        integer                                                  :: i, j, owner
+
+#ifdef MFC_MPI
+        integer :: ierr  !< Generic flag used to identify and report MPI errors
+#endif
+
+        ! Each round, the lowest rank holding an unlisted name broadcasts it
+        num_names = 0
+        do
+            owner = num_procs
+            do i = 1, nvtx_num_timers
+                if (f_nvtx_find(names(1:num_names), nvtx_timer_names(i)) == 0) then
+                    owner = proc_rank
+                    exit
+                end if
+            end do
+#ifdef MFC_MPI
+            call MPI_ALLREDUCE(MPI_IN_PLACE, owner, 1, MPI_INTEGER, MPI_MIN, MPI_COMM_WORLD, ierr)
+#endif
+            if (owner == num_procs .or. num_names == nvtx_max_timers) exit
+            num_names = num_names + 1
+            if (proc_rank == owner) names(num_names) = nvtx_timer_names(i)
+#ifdef MFC_MPI
+            call MPI_BCAST(names(num_names), nvtx_name_len, MPI_CHARACTER, owner, MPI_COMM_WORLD, ierr)
+#endif
+        end do
+
+        sec = f_nvtx_timer_seconds()
+        t_loc = 0._8
+        calls_loc = 0_8
+        do i = 1, num_names
+            j = f_nvtx_find(nvtx_timer_names(1:nvtx_num_timers), names(i))
+            if (j == 0) cycle
+            t_loc(i) = sec(j)
+            calls_loc(i) = nvtx_timer_calls(j)
+        end do
+
+#ifdef MFC_MPI
+        call MPI_REDUCE(t_loc, t_min, num_names, MPI_DOUBLE_PRECISION, MPI_MIN, 0, MPI_COMM_WORLD, ierr)
+        call MPI_REDUCE(t_loc, t_sum, num_names, MPI_DOUBLE_PRECISION, MPI_SUM, 0, MPI_COMM_WORLD, ierr)
+        call MPI_REDUCE(t_loc, t_max, num_names, MPI_DOUBLE_PRECISION, MPI_MAX, 0, MPI_COMM_WORLD, ierr)
+        call MPI_REDUCE(calls_loc, calls_max, num_names, MPI_INTEGER8, MPI_MAX, 0, MPI_COMM_WORLD, ierr)
+#else
+        t_min = t_loc; t_sum = t_loc; t_max = t_loc; calls_max = calls_loc
+#endif
+
+    end subroutine s_mpi_reduce_nvtx_timers
+
     !> Print a case file error with the prohibited condition and message, then abort execution.
     impure subroutine s_prohibit_abort(condition, message)
 

src/common/m_nvtx.f90

@@ -14,6 +14,15 @@ module m_nvtx
 
     character(len=256), private :: tempName
 
+    !> Per-rank wall-clock accumulation of every named range, enabled by phase_timing_wrt
+    integer, parameter           :: nvtx_name_len = 64, nvtx_max_timers = 256
+    logical                      :: nvtx_timing = .false.
+    integer                      :: nvtx_num_timers = 0
+    character(len=nvtx_name_len) :: nvtx_timer_names(nvtx_max_timers)
+    integer(c_int64_t)           :: nvtx_timer_ticks(nvtx_max_timers) = 0, nvtx_timer_calls(nvtx_max_timers) = 0
+    integer, private             :: depth = 0, stack_id(64)
+    integer(c_int64_t), private  :: stack_tick(64)
+
     type, bind(C) :: nvtxEventAttributes
         integer(c_int16_t) :: version = 1
         integer(c_int16_t) :: size = 48  !
@@ -64,6 +73,8 @@ subroutine nvtxStartRange(name, id)
         integer, intent(in), optional             :: id
         type(nvtxEventAttributes)                 :: event
 
+        if (nvtx_timing) call s_push_timer(name)
+
 #if defined(MFC_GPU) && defined(__PGI)
         tempName = trim(name) // c_null_char
 
@@ -85,6 +96,65 @@ subroutine nvtxEndRange
         call nvtxRangePop
 #endif
 
+        if (nvtx_timing .and. depth > 0) call s_pop_timer()
+
     end subroutine nvtxEndRange
 
+    !> Index of name in list, or 0 if absent
+    pure integer function f_nvtx_find(list, name) result(idx)
+
+        character(len=*), intent(in) :: list(:), name
+
+        do idx = 1, size(list)
+            if (list(idx) == name) return
+        end do
+        idx = 0
+
+    end function f_nvtx_find
+
+    !> Open a timed range, registering its name on first use (id 0 = untracked)
+    subroutine s_push_timer(name)
+
+        character(len=*), intent(in) :: name
+        integer                      :: id
+
+        id = f_nvtx_find(nvtx_timer_names(1:nvtx_num_timers), name)
+        if (id == 0 .and. nvtx_num_timers < nvtx_max_timers) then
+            nvtx_num_timers = nvtx_num_timers + 1
+            nvtx_timer_names(nvtx_num_timers) = name
+            id = nvtx_num_timers
+        end if
+        depth = depth + 1
+        stack_id(depth) = id
+        call system_clock(stack_tick(depth))
+
+    end subroutine s_push_timer
+
+    !> Close the innermost timed range and accumulate its elapsed ticks
+    subroutine s_pop_timer()
+
+        integer(c_int64_t) :: tick
+        integer            :: id
+
+        call system_clock(tick)
+        id = stack_id(depth)
+        if (id > 0) then
+            nvtx_timer_ticks(id) = nvtx_timer_ticks(id) + tick - stack_tick(depth)
+            nvtx_timer_calls(id) = nvtx_timer_calls(id) + 1
+        end if
+        depth = depth - 1
+
+    end subroutine s_pop_timer
+
+    !> Accumulated seconds of every registered range
+    function f_nvtx_timer_seconds() result(sec)
+
+        real(c_double)     :: sec(nvtx_num_timers)
+        integer(c_int64_t) :: rate
+
+        call system_clock(count_rate=rate)
+        sec = real(nvtx_timer_ticks(1:nvtx_num_timers), c_double)/real(rate, c_double)
+
+    end function f_nvtx_timer_seconds
+
 end module m_nvtx

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_global_parameters.fpp

@@ -353,6 +353,7 @@ contains
 
         ! Logistics (sim-specific)
         run_time_info = .false.
+        phase_timing_wrt = .false.
         t_step_old = dflt_int
 
         ! Computational domain parameters (sim-specific)