#!/bin/bash -eu
#SBATCH -A CMB103
#SBATCH --partition batch
#SBATCH -t 1:00:00
#SBATCH --exclusive
#SBATCH --mail-user=seshu@slac.stanford.edu
#SBATCH --mail-type=ALL
#SBATCH -C nvme

module load cray-python
module swap $LMOD_FAMILY_PRGENV PrgEnv-amd
module load amd/5.1
module load cray-pmi
# module load ums/default
# module load ums023/default
# module load hpctoolkit/2023.08

DEV=/lustre/orion/cmb138/scratch/seshuy/legion_s3d_subrank/
MECHANISM=hept
KERNEL_PATH="$DEV/kernels/${MECHANISM}_kernels"
export LD_LIBRARY_PATH=$DEV/build/$MECHANISM:$KERNEL_PATH:$DEV/legion/language/build/lib:$LD_LIBRARY_PATH
export FI_MR_CACHE_MONITOR=memhooks
export FI_CXI_RX_MATCH_MODE=software
export GASNET_OFI_DEVICE_0=cxi2
export GASNET_OFI_DEVICE_1=cxi1
export GASNET_OFI_DEVICE_2=cxi3
export GASNET_OFI_DEVICE_3=cxi0
export GASNET_OFI_DEVICE_TYPE=Node
export GASNET_OFI_RECEIVE_BUFF_SIZE=8M
export MPICH_OFI_NIC_POLICY=BLOCK

num_nodes=$SLURM_JOB_NUM_NODES
num_ranks=$SLURM_JOB_NUM_NODES

ranks_x=( [1]=1 [2]=2 [4]=2 [8]=4 [16]=4 [32]=8 [64]=8 [128]=16 [256]=16 [512]=32 [1024]=32 [2048]=64 [4096]=64 [8192]=128)
ranks_y=( [1]=1 [2]=1 [4]=2 [8]=2 [16]=4 [32]=4 [64]=8 [128]=8  [256]=16 [512]=16 [1024]=32 [2048]=32 [4096]=64 [8192]=64)

nx=${ranks_x[$num_nodes]}
ny=${ranks_y[$num_nodes]}

export RHSF_MAPMODE=allgpu
export GASNET_FREEZE_ON_ERROR=0
export RHSF_FUSE=200
export RHST_TRACE=1
export GASNET_BACKTRACE=1
export RHST_ARGS="-ll:gpu 8 -ll:cpu 1 -ll:util 3 -ll:bgwork 3 -ll:fsize 49152 -ll:csize 49152 -ll:zsize 1024 -ll:stacksize 16 -hip:hostreg 0 -hl:safe_mapper 1 -logfile run_%.log -lg:prof $num_ranks -lg:prof_logfile prof_%.gz"
#-lg:serializer ascii -lg:prof_footprint 0 -level legion_prof=2 -logfile log_%.prof -ll:force_kthreads 1

OUTDIR=./pwave_x_"$num_nodes"_"$MECHANISM"
mkdir -p $OUTDIR
scontrol write batch_script $SLURM_JOB_ID $OUTDIR/job.sh

./ammonia_s3d.py  \
    -x $DEV/fortran/$MECHANISM/s3d.x -l $DEV/build/$MECHANISM/librhst.so \
    -g 64x64x1024 -v 1x1x1 -t 40 -b $DEV/s3d/ -p "$nx"x"$ny"x1 -s slurm \
    --bomb-file=0d -k $OUTDIR -o $OUTDIR/ammonia.txt --title pressure_wave_test

# export LD_LIBRARY_PATH=/mnt/bb/seshuy/:$LD_LIBRARY_PATH

# sbcast --send-libs -pf $DEV/fortran/$MECHANISM/s3d.x /mnt/bb/seshuy/s3d.x
# sbcast -pf $DEV/build/$MECHANISM/librhst.so /mnt/bb/seshuy/librhsf.so

# for f in $DEV/build/$MECHANISM/*.so
# do
#     base="$(basename $f)"
#     sbcast -pf $f /mnt/bb/seshuy/$base
# done

# sbcast -pf $KERNEL_PATH/libsinge_hept_cpu.so /mnt/bb/seshuy/libsinge_hept_cpu.so
# sbcast -pf $KERNEL_PATH/libsinge_hept_gpu.so /mnt/bb/seshuy/libsinge_hept_gpu.so

# for f in $DEV/legion/language/build/lib/*
# do
#     base="$(basename $f)"
#     sbcast -pf $f /mnt/bb/seshuy/$base
# done

# ./ammonia_s3d.py  \
#     -x /mnt/bb/seshuy/s3d.x \
#     -g 64x64x1024 -v 1x1x1 -t 1 -b $DEV/s3d/ -p "$nx"x"$ny"x1 -s slurm \
#     --bomb-file=0d -k $OUTDIR -o $OUTDIR/ammonia.txt --title pressure_wave_test

srun -N $num_nodes -n $num_nodes --ntasks-per-node 1 ./copy_logs.py $OUTDIR/run

srun --export=NONE -N $num_nodes -n $num_nodes --ntasks-per-node 1 $DEV/Ammonia_Cases/dmesg.sh $OUTDIR

cp slurm-$SLURM_JOB_ID.out $OUTDIR
