# 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

###############################################################################
#                    Options ohh! too main options
###############################################################################
# This file organizes the plenty of options that once can pass to the profiler
# tool scripts for generating, compiling, verifying, and profiling IREE-compiled
# MLIR operations.
#
# The options are organized into groups: typical, compilation, iree-compile,
# verification, profiling, performance-reporting. Note that there is a function
# of each group.
###############################################################################


def add_typical_arguments(parser):
    """Adds typical command line arguments to the parser."""
    parser.add_argument(
        "--iree-bin-dir",
        default="./tools",
        help="Directory containing IREE binaries, "
        "e.g. iree-compile, iree-benchmark-module, "
        "iree-run-module",
    )
    parser.add_argument(
        "--generated-dir",
        default=".",
        help="The dispatch profiler scripts generate "
        "mlir dispatches, compiled vmfbs, and reference_chache "
        "containing golden npy files in the generated-dir",
    )
    parser.add_argument(
        "--operation-kind",
        "--op-kind",
        dest="operation_kind",
        default="all",
        help="Specifies the operation kinds to generate.",
        choices=["matmul", "conv2d", "all"],
    )
    parser.add_argument(
        "--dispatches",
        default="",
        help="Comma delimited list to filter dispatches by name. "
        "A dispatch is a combination of operation and tuning "
        "configuration.",
    )
    parser.add_argument(
        "--mlir-dialect",
        default="linalg",
        help="MLIR dialect entry point at which operation is emitter.",
        choices=["linalg"],
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Prints verbose output and commands executed.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Prints commands that will be executed without actually "
        "executing them.",
    )
    parser.add_argument(
        "--default-config",
        action="store_true",
        help="Adds a dispatch without a pre-defined "
        "tuning configuration. This dispatch will use "
        "default configuration from KernelsConfig.cpp.",
    )


def add_compilation_arguments(parser):
    """Adds compilation (not part of iree-compile) command line arguments to the parser."""
    compilation_parser = parser.add_argument_group(
        "Compilation", "Compilation related options."
    )
    compilation_parser.add_argument(
        "--num-cpu",
        "-j",
        dest="num_cpu",
        type=int,
        default=-1,
        help="Number of cpu threads to use for compilation.",
    )
    compilation_parser.add_argument(
        "--force-compile",
        action="store_true",
        help="Force re-compilation of the operation even " "if .vmfb file is present.",
    )


def add_iree_compile_arguments(parser):
    """Adds iree-compile command line arguments to the parser."""
    iree_compile_parser = parser.add_argument_group(
        "iree-compile", "iree-compile related options."
    )

    iree_compile_parser.add_argument(
        "--iree-hal-target-backends",
        "--device",
        dest="device",
        default="cuda",
        help="Target backends for executable compilation. ",
        choices=["cuda", "vulkan", "cpu"],
    )
    iree_compile_parser.add_argument(
        "--iree-hal-cuda-llvm-target-arch",
        "--cuda-arch",
        dest="cuda_arch",
        default="sm_80",
        help="Target architecture for the CUDA backend. ",
        choices=["sm_50", "sm_60", "sm_75", "sm_80", "sm_86"],
    )
    iree_compile_parser.add_argument(
        "--iree-hal-benchmark-dispatch-repeat-count",
        "--batch-size",
        dest="batch_size",
        default=100,
        help="Number of times dispatch is launched in a loop to "
        "amortize the launch overhead. This argument is used for "
        "iree-compile and iree-benchamrk-module. The value used by "
        "iree-compile and iree-benchamrk-module should be the same.",
    )
    iree_compile_parser.add_argument(
        "--iree-flow-split-matmul-reduction",
        "--split-k-slices",
        dest="split_k_slices",
        default="",
        help="Number of slices to split the reduction K-dimension.",
    )
    iree_compile_parser.add_argument(
        "--iree-codegen-llvmgpu-use-mma-sync",
        "--use-mma-sync",
        dest="use_mma_sync",
        action="store_true",
        help="Use mma.sync instructions.",
    )
    iree_compile_parser.add_argument(
        "--iree-codegen-llvmgpu-use-wmma",
        "--use-wmma",
        dest="use_wmma",
        action="store_true",
        help="Use wmma instructions.",
    )
    iree_compile_parser.add_argument(
        "--mlir-print-ir-after-all",
        "--print-ir-after-all",
        dest="mlir_print_ir_after_all",
        action="store_true",
        help="Prints IR after all transformations and dumps a "
        "file print_ir_after_*.mlir file.",
    )


def add_verification_arguments(parser):
    """Adds verification related arguments to the parser."""
    verification_parser = parser.add_argument_group(
        "Verification", "Verification related options."
    )

    verification_parser.add_argument(
        "--verification-enabled", default="True", type=str, help="Verify the operation."
    )
    verification_parser.add_argument(
        "--verification-providers",
        default="numpy",
        choices=["numpy"],
        help="Comma delimited list of verification providers.",
    )


def add_profiling_arguments(parser):
    """Adds profiling related arguments to the parser."""
    profiling_parser = parser.add_argument_group(
        "Profiling", "Profiling (iree-benchmark-module) related options."
    )

    profiling_parser.add_argument(
        "--profiling-enabled",
        "--benchmark",
        default="True",
        type=str,
        help="Benchmark the operation.",
    )
    profiling_parser.add_argument(
        "--benchmark-repetitions",
        default=5,
        type=int,
        help="Number of times benchmark is repeated "
        "and min, max, median, and average runtimes/gflops are "
        "reported.",
    )


def add_performance_report_arguments(parser):
    """Adds performance report related arguments to the parser."""

    performance_report_parser = parser.add_argument_group(
        "Performance Report", "Performance report related options."
    )

    performance_report_parser.add_argument(
        "--output", default="", help="Path to output file for csv readable results."
    )
    performance_report_parser.add_argument(
        "--append",
        action="store_true",
        help="Appends the results to existing file. "
        "o.w., the existing file is overwritten.",
    )
    performance_report_parser.add_argument(
        "--tags",
        default="",
        help="Inserts leading columns in output table "
        "and uniform values for each column. Useful for "
        "generating pivot tables.",
    )


def add_matmul_arguments(parser):
    """Adds matmul related arguments to the parser."""

    matmul_parser = parser.add_argument_group(
        "Matmul", "Matrix-multiplication related options."
    )
    matmul_parser.add_argument(
        "--problem-m",
        default="",
        help="M dimension of the matrix. "
        "--problem-m==<value>,<value_start:value_end:increment>*",
    )
    matmul_parser.add_argument(
        "--problem-n",
        default="",
        help="N dimension of the matrix."
        "--problem-n==<value>,<value_start:value_end:increment>*",
    )
    matmul_parser.add_argument(
        "--problem-k",
        default="",
        help="K dimension of the matrix."
        "--problem-k==<value>,<value_start:value_end:increment>*",
    )


###############################################################################
# Parser all the arguments for a script function:
# parse_generator_arguments() for generator.py
# parse_profiler_arguments() for profiler.py
###############################################################################


def parse_generator_arguments(parser):
    """Adds and parse all the arguments for the *generator.py* script."""
    add_typical_arguments(parser)
    add_matmul_arguments(parser)
    add_iree_compile_arguments(parser)
    args = parser.parse_args()
    return args


def parse_compile_arguments(parser):
    """Adds and parse all the arguments for the *compile.py* script."""
    add_typical_arguments(parser)
    add_compilation_arguments(parser)
    add_iree_compile_arguments(parser)
    args = parser.parse_args()
    return args


def parse_profiler_arguments(parser):
    """Adds and parse all the arguments for the *profiler.py* script."""
    add_typical_arguments(parser)
    add_compilation_arguments(parser)
    add_iree_compile_arguments(parser)
    add_verification_arguments(parser)
    add_profiling_arguments(parser)
    add_performance_report_arguments(parser)

    # Additional arguments for the profiler.
    parser.add_argument(
        "--save-cmds",
        action="store_true",
        help="Saves commands and their output that are executed "
        "by the profiler in a file.",
    )

    args = parser.parse_args()

    # Boolenize the string arguments from command line. For these args, it makes easier
    # to read and convey the meaning. The boolean arguments below are specified as:
    # `--argument=<true|false>`
    args.verification_enabled = (
        False if args.verification_enabled in ["False", "false", "0"] else True
    )

    args.profiling_enabled = (
        False if args.profiling_enabled in ["False", "false", "0"] else True
    )

    return args


###############################################################################
# Helper functions for parsing command line arguments.
###############################################################################
def get_cmd_line_argument_ranges(arg):
    """Returns a list of values generated by range of the form start:end:increment."""
    if not arg:
        return []
    if ":" not in arg:
        return [int(arg)]
    range_elements = arg.split(":")
    start = int(range_elements[0])
    end = int(range_elements[1])
    increment = int(range_elements[2]) if len(range_elements) == 3 else 1
    return range(start, end, increment)


def get_cmd_line_argument_list(arg):
    """Returns a list of values generated by comma delimited string."""
    values = arg.split(",")
    range_list = []
    for val in values:
        range_list += get_cmd_line_argument_ranges(val)
    return range_list