tune_models.py

import os, subprocess, argparse, time
import tune_ck as tc


def parse_args():
    parser = argparse.ArgumentParser(
        description="Tune CK GEMMs for one or more ONNX models")
    parser.add_argument('--models',
                        '-m',
                        nargs='+',
                        help='ONNX models to be tuned',
                        required=True)
    parser.add_argument('--batch_sizes',
                        '-b',
                        nargs='+',
                        help='Batch sizes to tune',
                        required=True)
    parser.add_argument('--sequence_length',
                        '-s',
                        type=int,
                        default=384,
                        help='Sequence length for transformer models')
    parser.add_argument('-n',
                        type=int,
                        default=16,
                        help='Number of instances to tune')
    args = parser.parse_args()
    return args


def tune_models(models, batch_sizes, seq_len, n):
    time_stamp = time.strftime("%Y_%m_%d_%H_%M")
    log_file = "ck_tuning_{}.log".format(time_stamp)
    json_file = "ck_tuning_{}.json".format(time_stamp)
    for model in models:
        for batch in batch_sizes:
            out = subprocess.run(
                'MIGRAPHX_LOG_CK_GEMM=1 ../build/bin/driver run {} -g --fill1 input_ids --input-dim @input_ids {} {}  | grep \'ck_gemm.*: \[{{\' | sort -u >> {}'
                .format(model, batch, seq_len, log_file),
                capture_output=True,
                check=True,
                shell=True)

    tc.tune(log_file, n, json_file)
    tuning_path = os.path.abspath(json_file)
    os.environ["MIGRAPHX_CK_TUNING"] = tuning_path
    print("\nTuning results have been saved to:\n{}\n".format(json_file))


def run(args):
    tune_models(args.models, args.batch_sizes, args.sequence_length, args.n)


run(parse_args())