xref: /aosp_15_r20/external/pytorch/torchgen/_autoheuristic/pad_mm/generate_heuristic_pad_mm.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3if [ $# -ne 1 ]; then
4    echo "Error: This script requires exactly one argument."
5    echo "`bash generate_heuristic_pad_mm.sh collect` to run benchmark and collect training data."
6    echo "`bash generate_heuristic_pad_mm.sh generate` to use the collected data to learn a heuristic."
7    exit 1
8fi
9
10MODE=$1
11
12# !!! SPECIFY THE GPUs THAT YOU WANT TO USE HERE !!!
13GPU_DEVICE_IDS="4,5"
14
15# !!! SPECIFY THE CONDA ENVIRONEMNT THAT YOU WANT TO BE ACTIVATED HERE !!!
16CONDA_ENV=heuristic-pr
17
18NUM_SAMPLES=2000
19
20# This is where AutoHeuristic will store autotuning results
21OUTPUT_DIR="a100"
22
23# !!! CHANGE THE NAME OF THE HEURISTIC IF YOU WANT TO LEARN A HEURISTIC FOR A GPU THAT IS NOT A100 !!!
24HEURISTIC_NAME="PadMMA100"
25
26BENCHMARK_SCRIPT="gen_data_pad_mm.py"
27
28TRAIN_SCRIPT="train_regression_pad_mm.py"
29
30bash ../generate_heuristic.sh ${MODE} ${GPU_DEVICE_IDS} ${CONDA_ENV} ${NUM_SAMPLES} ${OUTPUT_DIR} ${HEURISTIC_NAME} ${BENCHMARK_SCRIPT} ${TRAIN_SCRIPT}
31