#!/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 helpers for color printing.
source ../helpers.sh

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

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

# 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"

export OMP_NUM_THREADS=1
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"

# First instrument binary for faster profiling.
INST="$RESDIR/$BINARY.inst"
rm -rf omnitrace-ectrans-lam-benchmark-gpu-dp-acc.inst-output
srun --ntasks 1 --gpus 1 omnitrace-instrument -o $INST -- \
        ${INSTALLDIR}/${ECTRANS_DIR}/bin/$BINARY ${ARGS}

# Trace instrumented binary of ecTrans with given arguments.
srun --ntasks=$NTASKS --cpu-bind=${CPU_BIND} --output=$RESULTS/out.%j.%t \
        --error=$RESULTS/err.%j.%t --input=none \
        omnitrace-run -- "$SELECT_GPU_NAME" $INST ${ARGS}

# # Trace ecTrans with given arguments via runtime in.
# srun --ntasks=$NTASKS --cpu-bind=${CPU_BIND} --output=$RESULTS/out.%j.%t \
#         --error=$RESULTS/err.%j.%t --input=none \
#         omnitrace-run -- "$SELECT_GPU_NAME" ${INSTALLDIR}/${ECTRANS_DIR}/bin/$BINARY ${ARGS}

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

# Delete the intrumented binary
rm -rf "$INST"

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