# 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 *


################################################################################
class Dispatch:
  """
  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.
  """

  def __init__(self, operation, configuration):
    self.operation = operation
    self.configuration = configuration
    self.is_fused_dispatch = False

  def name(self):
    return f"{self.operation.name()}_{self.configuration.name()}"


################################################################################
class DispatchCollection:
  """
  DispatchCollection: A collection of dispatches that only vary in their 
    configurations but not in their operations. For example, a collection 
    of matmul dispatches with different tile sizes. 
    
    We can emit a single MLIR file for all the dispatches in a collection
    and compile with single run of iree-compile and them into a single executable
  """

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

  def get_dispatches(self):
    """Returns a list of dispatches in the collection."""
    dispatches = []
    for configuration in self.configuration_list:
      dispatches.append(Dispatch(self.operation, configuration))
    return dispatches

  def append(self, dispatch):
    """Appends a dispatch to the collection."""
    if dispatch.operation != self.operation:
      raise ValueError(
          f"operation {self.operation.name()} does not match the dispatch "
          f"collection operation name {dispatch.operation.name()}.")
    self.configuration_list.append(dispatch.configuration)

  def num_of_dispatches(self):
    """Returns number of dispatches in the collection."""
    return len(self.configuration_list)