1#!/bin/bash 2 3if [ $# -lt 8 ]; then 4 echo "Error: This script requires exactly at least 8 arguments." 5 exit 1 6fi 7 8MODE=$1 9GPU_DEVICE_IDS=$2 10CONDA_ENV=$3 11NUM_SAMPLES=$4 12OUTPUT_DIR=$5 13HEURISTIC_NAME=$6 14BENCHMARK_SCRIPT=$7 15TRAIN_SCRIPT=$8 16EXTRA_TRAIN_ARGS=$9 17 18mkdir -p ${OUTPUT_DIR} 19 20if [ "$MODE" = "collect" ]; then 21 # this will collect data for NUM_SAMPLES samples on the number of GPUs specified in GPU_DEVICE_IDS in parallel 22 bash ../collect_data.sh "python ${BENCHMARK_SCRIPT}" ${GPU_DEVICE_IDS} ${NUM_SAMPLES} ${CONDA_ENV} ${OUTPUT_DIR} 23elif [ "$MODE" = "generate" ]; then 24 # the bash script above generates one separate txt file per GPU 25 # if GPU_DEVICE_IDS=6,7, it will generate "data_6.txt", "data_7.txt" inside OUTPUT_DIR 26 # these files have to be merged into a single file before we can use AutoHeuristic to learn a heuristic 27 OUTPUT_FILE="${OUTPUT_DIR}/${HEURISTIC_NAME}.txt" 28 INPUT_FILES=$(echo $GPU_DEVICE_IDS | tr ',' '\n' | sed "s|^|${OUTPUT_DIR}/data_|" | sed 's/$/.txt/') 29 python ../merge_data.py ${OUTPUT_FILE} ${INPUT_FILES} 30 31 # This will learn a heuristic and generate the code into torch/_inductor/autoheuristic/artifacts/_${HEURISTIC_NAME}.py 32 python ${TRAIN_SCRIPT} ${OUTPUT_FILE} --heuristic-name ${HEURISTIC_NAME} ${EXTRA_TRAIN_ARGS} 33else 34 echo "Error: Invalid mode ${MODE}. Please use 'collect' or 'generate'." 35 exit 1 36fi 37