# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import argparse, os

from library import *
from manifest import *
from launchers import *
from concurrent.futures import ThreadPoolExecutor
from options import parse_compile_arguments

###############################################################################
# Compile main : The main entry point for the compile tool.
# This tool compiles IREE-compiled MLIR operations for a given backend device,
###############################################################################

if __name__ == "__main__":
    ###############################################################################
    # Parse command line arguments
    ###############################################################################
    parser = argparse.ArgumentParser(
        description="IREE Python compile tool for launching iree-compile for verification and "
        "profiling. Issues iree-compile for a given backend device and iree-compile "
        "flags. Uses ThreadPoolExecutor to launch multiple iree-compile processes "
        "in parallel."
    )

    args = parse_compile_arguments(parser)
    ###############################################################################

    # Manifests metadata for a group of accompanying operations and configurations.
    manifest = Manifest(args)
    manifest.load()

    # Try and use all CPUs to launch iree-compile in parallel.
    cpu_count = os.cpu_count()
    if args.num_cpu > 0:
        cpu_count = min(cpu_count, args.num_cpu)

    # For all the operations in the manifest, issue iree-compile for verification
    # and profiling in parallel using ThreadPoolExecutor and cpu_count threads.
    cmds = []
    with ThreadPoolExecutor(max_workers=cpu_count) as executor:
        # For all the operations in the manifest compile, verify, and profile.
        for _, dispatch_collection_list in manifest.dispatch_collection_map.items():
            for dispatch_collection in dispatch_collection_list:
                # Create an instance of operation_launcher.
                operation = dispatch_collection.operation
                operation_launcher = IreeToolsLauncher(args, operation)
                for configuration in dispatch_collection.configuration_list:
                    for compile_mode in [
                        CompilationMode.Profile,
                        CompilationMode.Verify,
                    ]:
                        cmds.append(
                            executor.submit(
                                operation_launcher.iree_compile, compile_mode
                            )
                        )

    # Wait for all the commands to complete.
    results = [cmd.result() for cmd in cmds]