# Copyright 2022 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
"""Utilities for handling the benchmark suite.

Benchmark artifacts should be generated by building the `iree-benchmark-suites`
CMake target, which put them in the following directory structure:

<root-build-dir>/benchmark_suites
└── <benchmark-category> (e.g., TFLite)
    ├── <benchmark-suite> (e.g., MobileBertSquad-fp32)
    │   ├── <benchmark-case> (e.g., iree-vulkan__GPU-Mali-Valhall__kernel-execution)
    │   │   ├── compilation_statistics.json
    │   │   ├── tool
    │   │   └── flagfile
    │   ├── ...
    │   │   ├── compilation_statistics.json
    │   │   ├── tool
    │   │   └── flagfile
    │   └── <benchmark_case>
    │   │   ├── compilation_statistics.json
    │       ├── tool
    │       └── flagfile
    └── vmfb
        ├── <compiled-iree-model>-<sha1>.vmfb
        ├── ...
        └── <compiled-iree-model>-<sha1>.vmfb
"""

import collections
import pathlib
import re

from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple
from common.benchmark_definition import IREE_DRIVERS_INFOS, DriverInfo
from e2e_test_artifacts import iree_artifacts
from e2e_test_framework.definitions import iree_definitions

# All benchmarks' relative path against root build directory.
BENCHMARK_SUITE_REL_PATH = "benchmark_suites"

MODEL_FLAGFILE_NAME = "flagfile"
MODEL_TOOLFILE_NAME = "tool"


@dataclass
class BenchmarkCase:
  """Represents a benchmark case.

    model_name: the source model, e.g., 'MobileSSD'.
    model_tags: the source model tags, e.g., ['f32'].
    bench_mode: the benchmark mode, e.g., '1-thread,big-core'.
    target_arch: the target CPU/GPU architature, e.g., 'GPU-Adreno'.
    driver_info: the IREE driver configuration.
    benchmark_tool_name: the benchmark tool, e.g., 'iree-benchmark-module'.
    benchmark_case_dir: the path to benchmark case directory.
    run_config: the run config from e2e test framework. This overrides the
      `benchmark_case_dir`.
  """

  model_name: str
  model_tags: Sequence[str]
  bench_mode: Sequence[str]
  target_arch: str
  driver_info: DriverInfo
  benchmark_tool_name: str
  benchmark_case_dir: pathlib.Path
  run_config: Optional[iree_definitions.E2EModelRunConfig] = None


# A map from execution config to driver info. This is temporary during migration
# before we can drop the DriverInfo.
EXECUTION_CONFIG_TO_DRIVER_INFO_KEY_MAP: Dict[Tuple[
    iree_definitions.RuntimeDriver, iree_definitions.RuntimeLoader], str] = {
        (iree_definitions.RuntimeDriver.LOCAL_TASK, iree_definitions.RuntimeLoader.EMBEDDED_ELF):
            "iree-llvm-cpu",
        (iree_definitions.RuntimeDriver.LOCAL_SYNC, iree_definitions.RuntimeLoader.EMBEDDED_ELF):
            "iree-llvm-cpu-sync",
        (iree_definitions.RuntimeDriver.LOCAL_TASK, iree_definitions.RuntimeLoader.VMVX_MODULE):
            "iree-vmvx",
        (iree_definitions.RuntimeDriver.LOCAL_SYNC, iree_definitions.RuntimeLoader.VMVX_MODULE):
            "iree-vmvx-sync",
        (iree_definitions.RuntimeDriver.VULKAN, iree_definitions.RuntimeLoader.NONE):
            "iree-vulkan",
        (iree_definitions.RuntimeDriver.CUDA, iree_definitions.RuntimeLoader.NONE):
            "iree-cuda",
    }


class BenchmarkSuite(object):
  """Represents the benchmarks in benchmark suite directory."""

  def __init__(self,
               suite_map: Dict[pathlib.Path, List[BenchmarkCase]],
               legacy_suite: bool = False):
    """Construct a benchmark suite.

    Args:
      suites: the map of benchmark cases keyed by category directories.
      legacy_suite: true if this is a legacy benchmark suite.
    """
    self.suite_map = suite_map
    self.category_map = dict((category_dir.name, category_dir)
                             for category_dir in self.suite_map.keys())
    self.legacy_suite = legacy_suite

  def list_categories(self) -> List[Tuple[str, pathlib.Path]]:
    """Returns all categories and their directories.

    Returns:
      A tuple of (category name, category dir).
    """
    category_list = [(name, path) for name, path in self.category_map.items()]
    # Fix the order of category list.
    category_list.sort(key=lambda category: category[0])
    return category_list

  # TODO(#11076): target_architectures should be a list of
  # common_definitions.DeviceArchitecture instead of string, after removing the
  # legacy path.
  def filter_benchmarks_for_category(
      self,
      category: str,
      available_drivers: Optional[Sequence[str]] = None,
      available_loaders: Optional[Sequence[str]] = None,
      target_architectures: Optional[Sequence[str]] = None,
      driver_filter: Optional[str] = None,
      mode_filter: Optional[str] = None,
      model_name_filter: Optional[str] = None) -> Sequence[BenchmarkCase]:
    """Filters benchmarks in a specific category for the given device.
      Args:
        category: the specific benchmark category.
        available_drivers: list of drivers supported by the tools. None means to
          match any driver.
        available_loaders: list of executable loaders supported by the tools.
          None means to match any loader.
        target_architectures: list of target architectures to be included. None
          means no filter.
        driver_filter: driver filter regex.
        mode_filter: benchmark mode regex.
        model_name_filter: model name regex.
      Returns:
        A list of matched benchmark cases.
    """

    category_dir = self.category_map.get(category)
    if category_dir is None:
      return []

    chosen_cases = []
    for benchmark_case in self.suite_map[category_dir]:
      driver_info = benchmark_case.driver_info

      driver_name = driver_info.driver_name
      matched_available_driver = (available_drivers is None or
                                  driver_name in available_drivers)
      matched_driver_filter = driver_filter is None or re.match(
          driver_filter, driver_name) is not None
      matched_driver = matched_available_driver and matched_driver_filter

      matched_loader = not driver_info.loader_name or available_loaders is None or (
          driver_info.loader_name in available_loaders)

      target_arch = benchmark_case.target_arch.lower()
      if target_architectures is None:
        matched_arch = True
      else:
        matched_arch = target_arch in target_architectures

      bench_mode = ','.join(benchmark_case.bench_mode)
      matched_mode = (mode_filter is None or
                      re.match(mode_filter, bench_mode) is not None)

      model_name_with_tags = benchmark_case.model_name
      if len(benchmark_case.model_tags) > 0:
        model_name_with_tags += f"-{','.join(benchmark_case.model_tags)}"
      if self.legacy_suite:
        # For backward compatibility, model_name_filter matches against the string:
        #   <model name with tags>/<benchmark case name>
        model_and_case_name = f"{model_name_with_tags}/{benchmark_case.benchmark_case_dir.name}"
      else:
        # For the new run option, we drop the obscure old semantic and only
        # search on model name and its tags.
        model_and_case_name = model_name_with_tags
      matched_model_name = (model_name_filter is None or re.match(
          model_name_filter, model_and_case_name) is not None)

      if (matched_driver and matched_loader and matched_arch and
          matched_model_name and matched_mode):
        chosen_cases.append(benchmark_case)

    return chosen_cases

  @staticmethod
  def load_from_run_configs(
      run_configs: Sequence[iree_definitions.E2EModelRunConfig],
      root_benchmark_dir: pathlib.Path):
    """Loads the benchmarks from the run configs.

    Args:
      run_configs: list of benchmark run configs.
    Returns:
      A benchmark suite.
    """

    suite_map = collections.defaultdict(list)
    for run_config in run_configs:
      module_gen_config = run_config.module_generation_config
      module_exec_config = run_config.module_execution_config
      target_device_spec = run_config.target_device_spec

      driver_info_key = EXECUTION_CONFIG_TO_DRIVER_INFO_KEY_MAP.get(
          (module_exec_config.driver, module_exec_config.loader))
      if driver_info_key is None:
        raise ValueError(
            f"Can't map execution config to driver info: {module_exec_config}.")
      driver_info = IREE_DRIVERS_INFOS[driver_info_key]

      target_arch = str(target_device_spec.architecture)
      model = module_gen_config.imported_model.model

      module_dir_path = iree_artifacts.get_module_dir_path(
          module_generation_config=module_gen_config,
          root_path=root_benchmark_dir)
      module_dir_path = pathlib.Path(module_dir_path)

      benchmark_case = BenchmarkCase(model_name=model.name,
                                     model_tags=model.tags,
                                     bench_mode=module_exec_config.tags,
                                     target_arch=target_arch,
                                     driver_info=driver_info,
                                     benchmark_tool_name=run_config.tool.value,
                                     benchmark_case_dir=module_dir_path,
                                     run_config=run_config)
      category = pathlib.Path(model.source_type.value)
      suite_map[category].append(benchmark_case)

    return BenchmarkSuite(suite_map=suite_map)

  @staticmethod
  def load_from_benchmark_suite_dir(benchmark_suite_dir: pathlib.Path):
    """Scans and loads the benchmarks under the directory."""

    suite_map: Dict[pathlib.Path,
                    List[BenchmarkCase]] = collections.defaultdict(list)
    for benchmark_case_dir in benchmark_suite_dir.glob("**"):
      model_dir = benchmark_case_dir.parent
      benchmark_name = benchmark_case_dir.name
      # Take the benchmark directory name and see if it matches the benchmark
      # naming convention:
      #   <iree-driver>__<target-architecture>__<benchmark_mode>
      segments = benchmark_name.split("__")
      if len(segments) != 3 or not segments[0].startswith("iree-"):
        continue

      config, target_arch, bench_mode = segments
      bench_mode = bench_mode.split(",")

      # The path of model_dir is expected to be:
      #   <benchmark_suite_dir>/<category>/<model_name>-<model_tags>
      category_dir = model_dir.parent
      model_name_with_tags = model_dir.name
      model_name_parts = model_name_with_tags.split("-", 1)
      model_name = model_name_parts[0]
      if len(model_name_parts) == 2:
        model_tags = model_name_parts[1].split(",")
      else:
        model_tags = []

      tool_name = (benchmark_case_dir / MODEL_TOOLFILE_NAME).read_text().strip()

      suite_map[category_dir].append(
          BenchmarkCase(model_name=model_name,
                        model_tags=model_tags,
                        bench_mode=bench_mode,
                        target_arch=target_arch,
                        driver_info=IREE_DRIVERS_INFOS[config.lower()],
                        benchmark_case_dir=benchmark_case_dir,
                        benchmark_tool_name=tool_name))

    return BenchmarkSuite(suite_map=suite_map, legacy_suite=True)