#!/bin/bash
# Copyright 2020 Joshua Bakita
# A unified script for timing:
# - SMT pairs with synchronous launches (hard real-time methodolgy)
# - SMT pairs with asynchronous launches (soft real-time methodolgy)
# - Unpaired tasks
# - The cache alloc vs WSS setting exploration for DIS
# These can all be done with:
# - No cache management
# - L3 cache splitting
# - L2+L3 cache splitting (if OS supported)

if grep -q "#define LITMUS 1" ../extra.h ./extra.h 2> /dev/null; then
    LITMUS=1
fi

if [ $# -lt 5 ]; then
    echo "Usage $0 -m MODE -p CPU -l LOOPS -b FILE [-B] [-I] RUN_ID"
    echo " -m base|dis|pair  Which benchmarking mode to use"
    echo " -p CPU       Which CPU to run each benchmark on"
    echo " -l LOOPS     How many loops of each benchmark to do"
    echo " -b FILE      List of benchmarks to execute. Optional tab-"
    echo "              delimited 2nd column specifies an input"
    echo "              generation command (fed to bench via stdin)"
    echo " -B           Enable background contenders on other CCXes"
    echo " -I xi|i3|i   Isolation mode: none (xi) (default), way-based"
    echo "              L3-only (i3), or color-based L2+L3 (i)"
    echo "Mode base requires no additional options."
    echo "Mode dis does not support the 2nd col. of -b and needs:"
    echo " -W FILE      List of working set sizes to pass gen_input.py"
    echo " -T FILE      Input template to pass to gen_input.py"
    echo " -C FILE      List of cache configurations to test"
    echo "Mode pair options:"
    echo " -P CPU       CPU to run the 2nd benchmark on"
    echo " -A           Async. (Don't synchronize the pair at job"
    echo "              boundries.) Typically used SMT in soft realtime"
fi

# Name options similarly to rtspin
while getopts “m:p:l:b:BI:W:T:C:P:A” opt; do
  case $opt in
# Required
    m) mode=$OPTARG ;;
    p) core=$OPTARG ;;
    l) maxJobs=$OPTARG ;;
    b) benchNames=$OPTARG ;;
# Optional
    B) contend=1 ;;
    I) iso_mode=$OPTARG ;;
# DIS
    W) wss_settings=$OPTARG ;;
    T) template_input=$OPTARG ;;
    C) cache_settings=$OPTARG ;;
# Pair
    P) core_two=$OPTARG ;;
    A) async=1 ;;
  esac
done
# Reset to read operands
shift $((OPTIND -1))
userRunID=$1

# Required argument checking
if [[ "$mode" != "base" && "$mode" != "dis" && $mode != pair ]]; then
    echo "Invalid argument: MODE must be either 'base' or 'dis'"
    exit
fi
if [[ ! -v core ]] || [[ ! -v maxJobs ]] || [[ ! -v benchNames ]]; then
    echo "Missing argument: -p, -l, and -b are required"
    exit
fi
if [[ ! -f "$benchNames" ]]; then
    echo "Invalid argument: $benchNames des not exist"
    exit
fi
if [[ ! -v userRunID ]]; then
    echo "Missing argument: RUN_ID is required"
    exit
fi

# DIS argument checking
if [[ "$mode" == "dis" ]] && [[ ! -v wss_settings || ! -v template_input || -z $cache_settings ]]; then
    echo "Missing argument: DIS needs -W FILE -T FILE and -C FILE"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$wss_settings" ]]; then
    echo "Invalid argument: $wss_settings does not exist"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$template_input" ]]; then
    echo "Invalid argument: $template_input does not exist"
    exit
fi
if [[ "$mode" == "dis" ]] && [[ ! -f "$cache_settings" ]]; then
    echo "Invalid argument: $cache_settings does not exist"
    exit
fi

# Pair argument checking
if [[ "$mode" == "pair" && ! -v core_two ]]; then
    echo "Missing argument: mode=pair reqires -P"
    exit
fi
if [[ "$mode" != "pair" ]] && [[ -v async || -v core_two ]]; then
    echo "Invalid argument: -A and -P require mode=pair"
    exit
fi

# Additional checks
if [[ $userRunID == "" ]]; then
    echo "Missing argument: a run ID is required"
    exit
fi
if [[ -v contend && ! -f "/playpen/mc2/imx6q-thrasher/thrasher" ]]; then
    echo "ERROR: cannot find thrasher binary for -B. Exiting..."
    exit
fi
if [[ $iso_mode == "i" && ! $(uname -a | grep "mc2") ]]; then
    echo "Isolation mode 'i' requires the MC^2 kernel. Exiting..."
    exit
fi

# Check permissions and update state
if [[ "$EUID" != 0 ]]; then
    echo "You need to be root to enable interrupt isolation and real-time execution!"
    exit
fi

echo "Loading benchmark names and input..."

# Read the names of each benchmark and load input.
# This might look a bit scary, but all it does is
# separate the two fields (1 or more tabs in-between),
# save the benchmark name, and execute the input
# command (storing its output in-memory for later).
j=0
IFS=$'\r\n'
while read i; do
    bench[$j]=$(echo $i | tr -s "\t" | cut -f1)
    input[$j]=$(eval $(echo $i | tr -s "\t" | cut -s -f2))
    j=$(( $j + 1 ))
done < $benchNames
echo "Making sure that binaries are up to date..."

for b in ${bench[@]}; do
    # Build "bin/bench" if bin directory
    if [[ -d bin ]]; then
        make bin/$b
    fi
    # If that failed, try a different name
    if [[ "$?" != 0 || ! -d bin ]]; then
        make $b
    fi
    # If bench is still not in bin or curr dir, fail
    if [[ ! -f "./$b" && ! -f "./bin/$b" ]]; then
        echo "Unable to find benchmark $b"
        exit
    fi
done
echo "Done. Disabling real-time throttling..."

# Turn off rt throttling
echo -1 > /proc/sys/kernel/sched_rt_runtime_us
echo "Done. Redirecting all interrupts to core 0..."

# Redirect all interrupts to core 0
i=0
for IRQ in /proc/irq/*
do
    # Skip default_smp_affinity
    if [ -d $IRQ ]; then
        irqList[$i]=$(cat $IRQ/smp_affinity_list)
        echo 0 2> /dev/null > $IRQ/smp_affinity_list
    fi
    i=$(( $i + 1 ))
done

echo "Done. Creating cleanup handler..."
function cleanup {
    # End contending tasks
    if [[ -v contend ]]; then
        killall thrasher
    fi

    # Restore cache access to everyone
    echo "L3:0=ffff;1=ffff;2=ffff;3=ffff" > /sys/fs/resctrl/schemata
    if [[ -e /sys/fs/resctrl/benchmarks ]]; then
        echo "" > /sys/fs/resctrl/benchmarks/cpus_list
    fi
    if [[ -e /sys/fs/resctrl/benchmarks2 ]]; then
        echo "" > /sys/fs/resctrl/benchmarks2/cpus_list
    fi

    # Put smp_affinty back the way it was
    i=0
    for IRQ in /proc/irq/*
    do
        if [ -d $IRQ ]; then
            echo ${irqList[$i]} 2> /dev/null > $IRQ/smp_affinity_list
        fi
        i=$(( $i + 1 ))
    done
}
# This sets the cleanup function to be called when we exit for any reason
trap cleanup EXIT

echo "Done. Checking for wbinvd module..."
if [[ ! -f "/proc/wbinvd" ]]; then
    echo "ERROR: wbinvd module not loaded. Exiting..."
    exit
fi

echo "Done. Setting cores to 'performance'..."
echo "performance" > /sys/devices/system/cpu/cpu$core/cpufreq/scaling_governor
if [[ -v core_two ]]; then
    echo "performance" > /sys/devices/system/cpu/cpu$core_two/cpufreq/scaling_governor
fi

# Enable L3 isolation
echo "Done. Setting up L3 isolation..."
mount -t resctrl resctrl /sys/fs/resctrl
# Reset global bandwith control and remove L3 from global
echo "L3:0=ffff;1=ffff;2=ffff;3=0000" > /sys/fs/resctrl/schemata
echo "MB:0=2048;1=2048;2=2048;3=2048" > /sys/fs/resctrl/schemata
# Alloc L3 to benchmark
# (this logic is safe when $core_two is empty)
if [[ $iso_mode == "i" ]]; then
    # With "i" we rely on numactl to do isolation
    # Full L3 access from both cores
    mkdir -p /sys/fs/resctrl/benchmarks
    # Must write both at the same time. Appends unsupported!
    echo $core,$core_two > /sys/fs/resctrl/benchmarks/cpus_list
    echo "L3:0=0000;1=0000;2=0000;3=ffff" > /sys/fs/resctrl/benchmarks/schemata
    # But each core is bound to its own numa mem node
    numa_arg0="--membind=0"
    numa_arg1="--membind=1"
elif [[ $iso_mode == "i3" ]]; then
    # With "i3" we rely on resctrl to do isolation
    mkdir -p /sys/fs/resctrl/benchmarks
    mkdir -p /sys/fs/resctrl/benchmarks2
    echo $core > /sys/fs/resctrl/benchmarks/cpus_list
    echo $core_two > /sys/fs/resctrl/benchmarks2/cpus_list
    echo "L3:0=0000;1=0000;2=0000;3=ff00" > /sys/fs/resctrl/benchmarks/schemata
    echo "L3:0=0000;1=0000;2=0000;3=00ff" > /sys/fs/resctrl/benchmarks2/schemata
    # Disable numa-based isolation
    numa_arg0="--interleave=all"
    numa_arg1="--interleave=all"
else
    iso_mode="xi"
    # No isolation via numactl or resctrl
    mkdir -p /sys/fs/resctrl/benchmarks
    echo $core,$core_two > /sys/fs/resctrl/benchmarks/cpus_list
    echo "L3:0=0000;1=0000;2=0000;3=ffff" > /sys/fs/resctrl/benchmarks/schemata
    # Disable numa-based isolation
    numa_arg0="--interleave=all"
    numa_arg1="--interleave=all"
fi
echo "Done. Verifying configuration with user..."

# Generate file name string
# We append to this as we parse the environment settings
runID=$(date +"%b%d-%H")

# Confirm configuration with user
echo "=== Global Config ==="
cat /sys/fs/resctrl/schemata
echo "=== NUMA Config ==="
echo "    '$numa_arg0' and '$numa_arg1'"
echo "=== Core $(cat /sys/fs/resctrl/benchmarks/cpus_list) Config ==="
cat /sys/fs/resctrl/benchmarks/schemata
if [[ -v core_two && $(cat /sys/fs/resctrl/benchmarks2/cpus_list) ]]; then
    echo "=== Core $(cat /sys/fs/resctrl/benchmarks2/cpus_list) Config ==="
    cat /sys/fs/resctrl/benchmarks2/schemata
fi

if [[ -v contend ]]; then
    echo "Will run 6 contending tasks"
    runID=$runID-c
else
    runID=$runID-xc
fi

runID=$runID-$iso_mode

if [[ -v async ]]; then
    echo "Will use asynchronous job pairing"
    runID=$runID-async
fi

if [[ $mode == "pair" && ! -v async ]]; then
    echo "Results will be saved as $runID-$userRunID-A.txt and $runID-$userRunID-B.txt"
else
    echo "Results will be saved as $runID-$userRunID.txt"
fi

if [[ $core == 0 || "$core_two" == 0 ]]; then
    echo "!!!!! DANGER !!!!!"
    echo "! Running real-time tasks on core 0 will conflict with interrupt"
    echo "! handling and likely freeze your system! Only continue if you know"
    echo "! exactly what you're doing!!"
    echo "!!!!!!!!!!!!!!!!!!"
fi
echo "Press enter to confirm environment, Ctrl-C to exit..."
read

# Start contending tasks
if [[ -v contend ]]; then
    echo "Done. Starting 6 contending tasks..."
    # Run two contending tasks on each other CCX
    taskset -c 1 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 2 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 5 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 6 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 9 /playpen/mc2/imx6q-thrasher/thrasher &
    taskset -c 10 /playpen/mc2/imx6q-thrasher/thrasher &
fi
sleep 1 # Wait for contending tasks to start
echo "Done. Beginning benchmarks..."

# Output coloring
FAIL_COLOR="\033[0;31m"
GOOD_COLOR="\033[0;32m"
RESET_COLOR="\033[0m"

# For each benchmark
for (( i = 0; i < ${#bench[@]} ; i++ )); do
    # Search for the benchmark in this directory, than /bin
    if [[ -f ${bench[$i]} ]]; then
        prefix="."
    else
        prefix="./bin"
    fi
    if [[ "$mode" == "base" ]]; then
        # Just run the benchmark if TACLeBench
        # Check if we're using LITMUS^RT or not
        if [[ -v $LITMUS ]]; then
            echo "${input[$i]}" | numactl $numa_arg0 $prefix/${bench[$i]} ${bench[$i]} $maxJobs $core $runID-$userRunID 1
        else
            # Remember: Unpaired tasks always get access to all colors (so we --interleave all)
            echo "${input[$i]}" | chrt -r 97 numactl $numa_arg0 taskset -c $core $prefix/${bench[$i]} ${bench[$i]} $maxJobs $core $runID-$userRunID 1
        fi
    elif [[ "$mode" == "dis" ]]; then
        # Loop through each WSS and cache config if DIS
        while read j; do # For cache setting
            echo $j > /sys/fs/resctrl/benchmarks/schemata
            while read ii; do # For WSS setting
                if [[ -v $LITMUS ]]; then
                    ./gen_input.py ${bench[$i]} inputs/${bench[$i]^}/$template_input $ii | numactl $numa_arg0 $prefix/${bench[$i]} $bench[$i]}-$ii-$j $maxJobs $core $runID 1
                else
                    ./gen_input.py ${bench[$i]} inputs/${bench[$i]^}/$template_input $ii | chrt -r 97 numactl $numa_arg0 taskset -c $core $prefix/${bench[$i]} ${bench[$i]}-$ii-$j $maxJobs $core $runID 1
                fi
            done < $wss_settings
        done < $cache_settings
    elif [[ "$mode" == "pair" ]]; then
        # Only async truly needs to do all-pairs
        if [[ ! -v async ]]; then
            start=$i
        else
            start=0
        fi
        for (( j = $start; j < ${#bench[@]} ; j++ )); do
            if [[ ! -v async ]]; then
                # Synchronize between pairs - original hard real-time SMT approach
                if [[ -v $LITMUS ]]; then
                    echo "${input[$i]}" | numactl $numa_arg0 taskset -c $core $prefix/${bench[$i]} ${bench[$i]} $maxJobs $core $core_two ${bench[$j]} $runID-$userRunID-A 1 & PID1=$!;
                    echo "${input[$j]}" | numactl $numa_arg1 taskset -c $core_two $prefix/${bench[$j]} ${bench[$j]} $maxJobs $core_two $core ${bench[$i]} $runID-$userRunID-B 1 & PID2=$!;
                else
                    echo "${input[$i]}" | chrt -r 97 numactl $numa_arg0 taskset -c $core $prefix/${bench[$i]} ${bench[$i]} $maxJobs $core $core_two ${bench[$j]} $runID-$userRunID-A 1 & PID1=$!;
                    echo "${input[$j]}" | chrt -r 97 numactl $numa_arg1 taskset -c $core_two $prefix/${bench[$j]} ${bench[$j]} $maxJobs $core_two $core ${bench[$i]} $runID-$userRunID-B 1 & PID2=$!;
                fi
                # We launched them asynchronously, so we have to wait
                wait $PID1 $PID2
            else
                # No synchronization between pairs - original soft real-time SMT approach
                if [[ -v $LITMUS ]]; then
                    echo "${input[$j]}" | numactl $numa_arg0 taskset -c $core $prefix/${bench[$j]} ${bench[$j]} -1 $core NULL 0 & PID1=$!;
                    echo "${input[$i]}" | numactl $numa_arg1 taskset -c $core_two $prefix/${bench[$i]} ${bench[$i]}"+"${bench[$j]} $maxJobs $core_two $runID-$userRunID 1 & PID2=$!;
                else
                    echo "${input[$j]}" | chrt -r 97 numactl $numa_arg0 taskset -c $core $prefix/${bench[$j]} ${bench[$j]} -1 $core NULL 0 & PID1=$!;
                    echo "${input[$i]}" | chrt -r 97 numactl $numa_arg1 taskset -c $core_two $prefix/${bench[$i]} ${bench[$i]}"+"${bench[$j]} $maxJobs $core_two $runID-$userRunID 1 & PID2=$!;
                fi
                wait $PID2
                kill $PID1
            fi
            if [[ $? != 0 ]]; then
                echo -e ${FAIL_COLOR}FAILED${RESET_COLOR}: ${bench[$i]} ${bench[$j]}
            else
                echo -e ${GOOD_COLOR}COMPLETE${RESET_COLOR}: ${bench[$i]} ${bench[$j]}
            fi
        done
    fi
    if [[ $? != 0 ]]; then
        echo -e ${FAIL_COLOR}FAILED${RESET_COLOR}: ${bench[$i]}
    else
        echo -e ${GOOD_COLOR}COMPLETE${RESET_COLOR}: ${bench[$i]}
    fi
done