#!/usr/bin/env bash
# ------------------------------------------------------------------------------
# This file runs the ecTrans dwarf and profiles the execution.
# The script is meant to support submitted jobs and interactive nodes.
# ------------------------------------------------------------------------------

# Load directory structure of installation.
source ../../dirs.sh

# Load required modules.
source ${BASEDIR}/load_modules.sh

# Load helpers for color printing.
source ${BASEDIR}/helpers.sh

# Set binary, results directory name, and ntasks to ENV value or default.
[ -z "$BINARY" ] && BINARY="ectrans-lam-benchmark-dp"
[ -z "$RESDIR" ] && RESDIR="profile_lumi-g"

# Set runtime arguments to ENV value or default.
[ -z "$NPROMA" ] && NPROMA=32
[ -z "$NFLD" ] && NFLD=1
[ -z "$NLEV" ] && NLEV=10
[ -z "$NITER" ] && NITER=10
[ -z "$NLAT" ] && NLAT=128
[ -z "$NLON" ] && NLON=128

# Specify where to store results, create folder and cd into it.
RESULTS="$RESDIR"
rm -rf $RESULTS
mkdir -p $RESULTS

# Create "select_gpu" script as indicated by Lumi documentation.
SELECT_GPU_NAME=./select_gpu-$RESDIR
rm -rf "$SELECT_GPU_NAME"
cat << EOF > "$SELECT_GPU_NAME"
#!/bin/bash

export ROCR_VISIBLE_DEVICES=\$SLURM_LOCALID
exec \$*
EOF
chmod +x "$SELECT_GPU_NAME"

# Setup GPU and OpenMP aware hardware config.
CPU_BIND="mask_cpu:7e000000000000,7e00000000000000"
CPU_BIND="${CPU_BIND},7e0000,7e000000"
CPU_BIND="${CPU_BIND},7e,7e00"
CPU_BIND="${CPU_BIND},7e00000000,7e0000000000"

# Set number of threads to 64 for CPU execution.
export OMP_NUM_THREADS=64
export MPICH_GPU_SUPPORT_ENABLED=1

# Set ecTrans runtime arguments.
ARGS="--nproma $NPROMA --vordiv --scders --uvders --nfld $NFLD --nlev $NLEV \
        --norms --niter $NITER --nlat $NLAT --nlon $NLON"
EXEC="$SELECT_GPU_NAME $BINARY $ARGS"
info "[$(date +'%H:%M:%S')]  Executing: $BINARY $ARGS"

# Trace ecTrans with given arguments.
ROCMFILES="$RESULTS/rocprof.csv"
srun --ntasks=1 --cpu-bind=${CPU_BIND} --output=$RESULTS/out.%j.%t \
        --error=$RESULTS/err.%j.%t --input=none \
        rocprof --hip-trace --roctx-trace -o $ROCMFILES \
        "$SELECT_GPU_NAME" ${INSTALLDIR}/${ECTRANS_DIR}/bin/$BINARY ${ARGS}

# Acquire HWCs of ecTrans with given arguments.
# Set the dir of this script in order to link counters correctly.
PROFILING=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
srun --ntasks=1 --cpu-bind=${CPU_BIND} --output=$RESULTS/out.%j.%t \
        --error=$RESULTS/err.%j.%t --input=none \
        rocprof -i "$PROFILING/rocprof_counters.txt" -o $ROCMFILES \
        "$SELECT_GPU_NAME" ${INSTALLDIR}/${ECTRANS_DIR}/bin/$BINARY ${ARGS}

# Delete the "select_gpu_sbatch" script.
rm -rf "$SELECT_GPU_NAME"

# Output succesfull run.
success "[$(date +'%H:%M:%S')]  Finished the profiling run of ecTrans."
