#!/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.
rm -rf $RESDIR
mkdir -p $RESDIR
cd $RESDIR

# Specify where to store slurm output and create folder.
RESULTS="slurm_out"
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=${SLURM_CPUS_PER_TASK}
export MPICH_GPU_SUPPORT_ENABLED=1

# Setup Extrae.
EXTRAE_BASEDIR="/project/project_465000527/perf_tools/extrae/extrae_user_fx"
source $EXTRAE_BASEDIR/env.sh
# Set ecTrans runtime arguments.
ARGS="--nproma $NPROMA --vordiv --scders --uvders --nfld $NFLD --nlev $NLEV \
        --norms --niter $NITER --nlat $NLAT --nlon $NLON"
info "[$(date +'%H:%M:%S')]  Executing: $BINARY $ARGS"

# Trace ecTrans with given arguments.
srun --cpu-bind=${CPU_BIND} --output=$RESULTS/out.%j.%t \
        --error=$RESULTS/err.%j.%t --input=none \
        ../trace.sh "$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."
