#!/bin/bash -l
#SBATCH --job-name=allreduce-bench
#SBATCH --nodes=36
#SBATCH --ntasks-per-node=8
#SBATCH --cpus-per-task=7
#SBATCH --gpus-per-node=8
#SBATCH --time=0:20:00
#SBATCH --account=phy157_dwf
#SBATCH --gpu-bind=none
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH -S 0

##############################################################################
# Benchmark_allreduce at 1, 4 and 36 nodes from one allocation: MPI_Allreduce
# (host and device buffers) vs CartesianRingAllReduce vs flat RingAllReduce
# (device, and host with explicit staging), 4 KB .. 512 MB, RealF and RealD.
# Decides whether GlobalSumVector wants a size threshold or simply the ring
# (deterministic, no size cliff) everywhere.  Cray MPICH's device-buffer
# allreduce aborted above ~8 MB (MPI_FLOAT) in earlier runs: MPI-dev is
# attempted only up to BENCH_MPI_DEV_MAX_MB (4); raise it to probe the cliff.
#
# Readout: the tables (one per type per node count).  Columns are ms (min of
# reps, slowest rank), payload GB/s and "wire" GB/s (bytes each ring actually
# moves / time).  Compare MPI-host with cart-dev / flat-dev at each size; the
# crossovers, if any, should move with P as sum_d 2(P_d-1) vs 2(P-1) steps.
##############################################################################

cat << EOF > select_gpu
#!/bin/bash
export GPU_MAP=(0 1 2 3 7 6 5 4)
export NUMA_MAP=(3 3 1 1 2 2 0 0)
export GPU=\${GPU_MAP[\$SLURM_LOCALID]}
export NUMA=\${NUMA_MAP[\$SLURM_LOCALID]}
export HIP_VISIBLE_DEVICES=\$GPU
unset ROCR_VISIBLE_DEVICES
if [ \$SLURM_PROCID = "0" ]; then echo \$*; fi
exec numactl -m \$NUMA -N \$NUMA \$*
EOF
chmod +x ./select_gpu

root=$HOME/ParallelIO/systems/Frontier
source $root/sourceme-rocm7.2.sh

export OMP_NUM_THREADS=7
export FI_MR_CACHE_MONITOR=kdreg2       # REQUIRED for device-buffer MPI on Slingshot: libfabric memhooks monitor (default) is defective, see systems/WorkArounds.txt (libfabric #11451)
export MPICH_GPU_SUPPORT_ENABLED=1
export MPICH_SMP_SINGLE_COPY_MODE=CMA
export MPICH_OFI_NIC_POLICY=GPU
module load libfabric

BIN=$root/benchmarks/Benchmark_allreduce
OPTS1="--accelerator-threads 8 --shm 4096 --shm-mpi 1 --device-mem 32000"
export BENCH_MIN_KB=4
export BENCH_MAX_MB=512
export BENCH_REPS=5
export BENCH_MPI_DEV_MAX_MB=4
export BENCH_SCALAR_REPS=200

run_cell () {
  name=$1; nodes=$2; ntask=$3; geom=$4; vol=$5
  echo "----- $name : $nodes nodes, $ntask ranks, --mpi $geom -----"
  fname=log.allreduce.$name
  srun -N$nodes -n$ntask --kill-on-bad-exit=1 ./select_gpu $BIN --mpi $geom --grid $vol $OPTS1 > $fname 2>&1
  echo "  exit $?"; sleep 20
  grep -h "====\|bytes  *elements\|| .* |\| min .* us" $fname | sed 's/^Grid : Message : [0-9.]* s : //' | cut -c1-175
}

#          name        nodes ntask  geom      vol
#run_cell   N1_8gcd     1     8      1.1.2.4   16.16.32.64      # done 2026-08-27
#run_cell   N4_32gcd    4     32     2.2.2.4   32.32.32.64      # done: MPI <1 MB, rings above, 2-2.5x at 512 MB
#run_cell   N36_288gcd  36    288    3.6.4.4   48.48.48.96      # done: MPI <2 MB, cart 2-48 MB, flat above (3x at 512 MB);
                                                               #       SCALAR: GlobalSumP2P 0.4 ms min / 12.8 ms MEAN vs MPI 45 us (!)

# SCALAR-ONLY sweep at 288 ranks (BENCH_MAX_MB=0 skips the vector tables; each cell ~30 s):
# S0 back-to-back as measured; S1 with ~2 ms of GPU work between sums (the solver's condition:
# fine-smoother steps are 16 ms apart incl. a ~10 ms matvec, so the solver cannot be paying
# 12 ms per reduction); S2-S3 one environment knob each.  Readout: the five "min/mean us" lines.
export BENCH_MAX_MB=0
scalar_cell () { name=$1; shift; run_cell_env "$name" 36 288 3.6.4.4 48.48.48.96 "$@"; }
run_cell_env () {
  name=$1; nodes=$2; ntask=$3; geom=$4; vol=$5; shift 5
  echo "----- $name : $* -----"
  fname=log.allreduce.$name
  env "$@" srun -N$nodes -n$ntask --kill-on-bad-exit=1 ./select_gpu $BIN --mpi $geom --grid $vol $OPTS1 > $fname 2>&1
  echo "  exit $?"; sleep 20
  grep -h "==== SCALAR\|gap kernel\| min .* us" $fname | sed 's/^Grid : Message : [0-9.]* s : //' | cut -c1-140
}
scalar_cell  S0_tight        BENCH_SCALAR_GAP_US=0
scalar_cell  S1_gap2ms       BENCH_SCALAR_GAP_US=2000
scalar_cell  S2_gap2ms_omp1  BENCH_SCALAR_GAP_US=2000 OMP_NUM_THREADS=1
scalar_cell  S3_gap2ms_nogpu BENCH_SCALAR_GAP_US=2000 MPICH_GPU_SUPPORT_ENABLED=0

echo "========================================================="
echo "done: see log.allreduce.N1_8gcd, N4_32gcd, N36_288gcd"
echo "========================================================="