optimize.py

import os
import argparse

from TuningDriver.interface import optimize_run, check_optimization_progress, merge, postprocess


def main(hipblaslt_path, lib_dir='lib', log_summary=None, devices="0,1,2,3,4,5,6,7", workdir="workdir"):
    
    input_dir = os.path.join(workdir, 'tunings')
    if not os.path.isdir(input_dir):
        raise ValueError(f"'tunings' directory not found in {workdir}")

    try:
        devices = list(set([int(d) for d in devices.split(',')]))
    except ValueError as e:
        raise ValueError(f'Error parsing devices {devices}')

    if len(devices) == 0:
        raise ValueError(f'Need at least 1 devices to run the optimization.')

    # 1 - OPTIMIZATION STEP
    os.makedirs( os.path.join(workdir,'lib'), exist_ok=True)
    
    n_gemms, n_completed = check_optimization_progress(input_dir)
    if n_gemms > n_completed:
        print(f'Running optimization for {n_gemms - n_completed} GEMMS. This will take a while...')
        optimize_run(input_dir, devices=devices)
        # TODO check progress

    if n_gemms == 0:
        raise ValueError(f"No GEMMs found in {input_dir}")
    
    _, n_completed = check_optimization_progress(input_dir)
     
    if n_completed == 0:
        raise ValueError(f"No GEMMs completed out of {n_gemms} in {input_dir}")
    
    print(f"Completed {n_completed}/{n_gemms} GEMMs")
    
    # 2 - MERGE KERNELS
    print(f'Merging new kernels into a single library...')
    merge(hipblaslt_path, input_dir,  os.path.join(workdir, 'lib'), workdir)

    # 3 - BENCHMARK AND WRITE FINAL LIBRARY
    print(f'Postprocessing new kernels...')
    postprocess(hipblaslt_path, 
                lib_dir,
                os.path.join(workdir, 'lib'), 
                os.path.join(workdir, 'tensile'), 
                log_summary=log_summary, 
                device=devices[0], 
                workdir=workdir)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=f"Merge tuning results into a single library.")
    parser.add_argument('hipblaslt_path', help='Path to hipBLASLt', type=str)
    parser.add_argument("--workdir", "-w", default="workdir", help="Dir to store intermediate files.")
    parser.add_argument('--library_dir', help='Final library output directory', type=str, default='final')
    parser.add_argument('--log_summary', help="CSV file generated by gen_configs.py script. Will be used to report the weighted uplift.", type=str)
    parser.add_argument("--devices", "-d", action="store", type=str, default="0,1,2,3,4,5,6,7",
                        help='Comma-separated list of device IDs to use. Ex: 0,3,4,5')
    args = parser.parse_args()

    main(args.hipblaslt_path,
         lib_dir=args.library_dir,
         log_summary=args.log_summary,
         devices=args.devices,
         workdir=args.workdir)