# 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

from library import *
from matmul import ReferenceMatmulOp
from batch_matmul import ReferenceBatchMatmulOp
from pathlib import Path
import subprocess


class IreeToolsLauncher:
    """Launcher for IREE tools."""

    def __init__(self, args, operation):
        self.operation = operation

        self.generated_path = Path(args.generated_dir, "generated", args.mlir_dialect)

        self.args = args
        self.benchmark_dispatch_repeat_count = args.batch_size
        self.batch_size = args.batch_size

        # paths to source dispatch mlir, compiled vmfb, and logs.
        self.operation_path = self.generated_path.joinpath(
            OperationKindNames[operation.operation_kind], operation.name()
        )

        self.source_mlir_file = self.operation_path.joinpath(
            operation.name()
        ).with_suffix(".mlir")

        # path to cached numpy refernece input and expected output files.
        self.op_reference_cache_path = Path(
            args.generated_dir, "generated", "reference_cache", operation.name()
        )

        if not self.op_reference_cache_path.exists():
            self.op_reference_cache_path.mkdir(parents=True, exist_ok=True)

        # path to iree-compile tool. (for compiling the input mlir file to vmfb)
        self.iree_compile_path = Path(args.iree_bin_dir, "iree-compile")

        # path to iree-benchmark-module tool. (for performance benchmarking and profiling)
        self.iree_benchmark_module_path = Path(
            args.iree_bin_dir, "iree-benchmark-module"
        )

        # path to iree-run-module tool. (for verification)
        self.iree_run_module_path = Path(args.iree_bin_dir, "iree-run-module")

        # output vmfb files for verification and profiling.
        vmfb_filename = f"{operation.name()}"

        if operation.operation_kind == OperationKind.SplitkMatmul:
            split_k_suffix = "_".join(["split_k_slice", str(operation.split_k_slices)])
            vmfb_filename = f"{vmfb_filename}_{split_k_suffix}"

        self.vmfb_verify_filepath = self.operation_path.joinpath(
            self.operation.name()
        ).with_name(f"{vmfb_filename}_verify.vmfb")
        self.vmfb_profile_filepath = self.operation_path.joinpath(
            self.operation.name()
        ).with_name(f"{vmfb_filename}_profile.vmfb")

        # reference implementation for the operation_kind.
        self.reference_impl_map = {
            OperationKind.Matmul: ReferenceMatmulOp,
            OperationKind.SplitkMatmul: ReferenceMatmulOp,
            OperationKind.BatchMatmul: ReferenceBatchMatmulOp,
        }

    def iree_compile(self, compilation_mode):
        """Compiles the input mlir file to vmfb file."""

        benchmark_dispatch_repeat_count = (
            self.benchmark_dispatch_repeat_count
            if compilation_mode == CompilationMode.Profile
            else 1
        )
        vmfb_filepath = (
            self.vmfb_profile_filepath
            if compilation_mode == CompilationMode.Profile
            else self.vmfb_verify_filepath
        )

        # Base iree-compile commandline
        cmd = [
            f"{self.iree_compile_path}",
            f"{self.source_mlir_file}",
            "-o",
            f"{vmfb_filepath}",
        ]

        # General compilation options
        cmd += [f"--iree-hal-target-backends={self.args.device}"]

        if self.args.device == "cuda":
            cmd += [f"--iree-hal-cuda-llvm-target-arch={self.args.cuda_arch}"]
        if self.operation.operation_kind == OperationKind.SplitkMatmul:
            cmd += [
                f"--iree-flow-split-matmul-reduction={self.operation.split_k_slices}"
            ]
        if self.args.use_mma_sync:
            cmd += [f"--iree-codegen-llvmgpu-use-mma-sync"]
        if self.args.use_wmma:
            cmd += [f"--iree-codegen-llvmgpu-use-wmma"]

        # Compilation options for profiling
        cmd += [
            f"--iree-hal-benchmark-dispatch-repeat-count={benchmark_dispatch_repeat_count}"
        ]

        # Appends print ir options at the end of the command line.
        if self.args.mlir_print_ir_after_all:
            cmd += [f"--mlir-print-ir-after-all"]

        if not vmfb_filepath.exists() or self.args.force_compile:
            complie_mode_str = CompilationModeNames[compilation_mode]

            print(f"[Compiling ({complie_mode_str})] {' '.join(cmd)}")

            iree_compile_stdout_filepath = self.operation_path.joinpath(
                "iree_compile_cmd_stdout.mlir"
            )

            with open(iree_compile_stdout_filepath, "w") as fp:
                subprocess.run(cmd, stderr=fp)

        elif self.args.verbose:
            print(
                f"Skipping compilation of operation: {vmfb_filepath} since it already exists."
            )

    def verify(self, configuration):
        """Verifies the operation with a given configuration."""
        # First compile the operation to a vmfb file.
        self.iree_compile(CompilationMode.Verify)

        # Verify using random data distribution.
        reference_run = self.reference_impl_map[self.operation.operation_kind](
            self.operation,
            self.op_reference_cache_path,
            Distribution.Random,
            Distribution.Random,
        )

        if not reference_run.is_cached():
            reference_run()

        # Commandline `iree-run-module` for verification.
        cmd = [
            f"{self.iree_run_module_path}",
            f"--module={self.vmfb_verify_filepath}",
            f"--device={self.args.device}",
        ]

        # Operation-specific verification command-line.
        cmd.append(f"--function={self.operation.name()}_{configuration.name()}")
        for input_file_path in reference_run.get_input_filepaths():
            cmd.append(f"--input=@{input_file_path}")
        for output_file_path in reference_run.get_output_filepaths():
            cmd.append(f"--expected_output=@{output_file_path}")

        # Print the command if verbose.
        if self.args.verbose:
            print(f"[Verification] {' '.join(cmd)}")

        # Launch verification.
        cmd_output = subprocess.check_output(cmd, text=True)

        # Save the verification command and the output, only if requested
        # (file writing could slow down the verification).
        if self.args.save_cmds:
            filepath = self.operation_path.joinpath("iree_run_module.stdout")
            with open(filepath, "w") as fp:
                fp.write(f"[Command] $ {' '.join(cmd)}\n")
                fp.write(cmd_output)

        # Parse the verification output.
        m = re.search(r"\[(?P<verification_result>[a-zA-Z]+)\]", cmd_output)
        if m is None:
            raise ValueError(
                f"Failed to parse verification output by iree-run-module: {cmd_output}"
            )
        verification_result = m.group("verification_result")

        if self.args.verbose or verification_result != "SUCCESS":
            print(cmd_output)

        return verification_result

    def profile(self, configuration):
        """Profiles the operation with a given configuration."""
        # First compile the operation to a vmfb file.
        self.iree_compile(CompilationMode.Profile)

        # Commandline `iree-benchmark-module` for profiling.
        cmd = [
            f"{self.iree_benchmark_module_path}",
            f"--module={self.vmfb_profile_filepath}",
            f"--device={self.args.device}",
        ]

        # Profiling specific flags.
        cmd += [f"--benchmark_repetitions={self.args.benchmark_repetitions}"]
        cmd += [f"--batch_size={self.batch_size}"]

        # Operation-specific profiling command-line.
        cmd += [f"--function={self.operation.name()}_{configuration.name()}"]
        cmd += [f"--input={self.operation.lhs_npy_shape()}"]
        cmd += [f"--input={self.operation.rhs_npy_shape()}"]

        # Print the command if verbose.
        if self.args.verbose:
            print(f"[Profiling] {' '.join(cmd)}")

        # Launch profiling.
        cmd_output = subprocess.check_output(cmd, text=True, stderr=subprocess.STDOUT)

        # Save the profiling command and the output, only if requested
        # (file writing could slow down the profiling).
        if self.args.save_cmds:
            filepath = self.operation_path.joinpath("iree_benchmark_module.stdout")
            with open(filepath, "w") as fp:
                fp.write(f"[Command] $ {' '.join(cmd)}\n")
                fp.write(cmd_output)

        # Parse the profiling output.
        m = re.search(r"real_time_median\s+(?P<runtime>\d+.\d+)\s+ms", cmd_output)
        if m is None:
            raise ValueError(
                f"Failed to parse runtime from benchmark result: {cmd_output}"
            )
        runtime_in_ms = float(m.group("runtime"))
        return runtime_in_ms