# 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

from library import *
from matmul import *
from batch_matmul import *
from manifest import *
from performance_report import *
from launchers import *
from options import parse_profiler_arguments

###############################################################################
# Profiler main : The main entry point for the profiler tool.
###############################################################################
# This tool compiles, verifies, and profiles IREE-compiled MLIR operations for
# a given backend device, compiler flags, and tuning configuration.
#
# The dispatch profiler tool is organized based on below defintions:
# Operation: A MLIR operation that is generated or consumed by the
#       dispatch_profiler. For example, linalg.matmul, linalg.conv2d, etc.
# Configuration: A set of compile parameters that are used by iree-compile the
#       to choose a compilation pipeline (e.g. LLVMGPUTensorCore,
#       LLVMGPUTensorCoreMmaSync, LLVGPUCPU, etc.), performance tuning parameters
#       (e.g. workgroup size, tile size etc.).
# Dispatch: A combination of an operation and a configuration is launched by the
#       dispatch profiler for verification and performance profiling. Note that
#       a dispatch is not a MLIR operation it is binary executable that is launched
#       by the profiler. Additionaly, the goal of the tool is to also profile the
#       performance of the fusions and a dispatch for fusion is a combination of
#       multiple operations glued together and compiled into a single dispatch.
###############################################################################

if __name__ == "__main__":
    ###############################################################################
    # Parse command line arguments
    ###############################################################################
    parser = argparse.ArgumentParser(
        description="IREE Python profiler tool for "
        "verifcation and performance profiling tool for IREE-compiled "
        "MLIR operations."
    )

    args = parse_profiler_arguments(parser)
    ###############################################################################

    # Create manifest object and load dispatches.
    manifest = Manifest(args)
    manifest.load()

    # Performance report
    perf_report = PerformanceReport(args)

    # For all the operations in the manifest compile (if needed), verify, and profile.
    for _, dispatch_collection_list in manifest.dispatch_collection_map.items():
        for dispatch_collection in dispatch_collection_list:
            operation = dispatch_collection.operation
            # Select and create an instance of operation_launcher for the operation.
            operation_launcher = IreeToolsLauncher(args, operation)
            for configuration in dispatch_collection.configuration_list:
                # Create a dispatch object.
                dispatch = Dispatch(operation, configuration)

                # Skip the dispatch if filter returns false.
                if not manifest.is_enabled(dispatch):
                    continue

                # If dry run is enabled, skip the dispatch.
                if args.dry_run:
                    print(f"[Dry run] : {dispatch.name()}")
                    continue

                # Initialize verification and profiling results.
                verification_result = (
                    "Not verified" if not args.verification_enabled else "Failed"
                )
                runtime = -1.0

                # Launch the operation dispatches for verification and profiling.
                if args.verification_enabled:
                    verification_result = operation_launcher.verify(configuration)
                if args.profiling_enabled:
                    runtime = operation_launcher.profile(configuration)

                # Create performance result.
                result = PerformanceResult(
                    operation, configuration, verification_result, runtime
                )

                # Print the performance result.
                result.print()

                # Append the performance result to the performance report.
                perf_report.append_perf_result(result)