text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you under the Apache License, Version 2.0 (the
  # "License"); you may not use this file except in compliance
  # with the License.  You may obtain a copy of the License at
  #
  #   http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing,
  # software distributed under the License is distributed on an
  # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  # KIND, either express or implied.  See the License for the
  # specific language governing permissions and limitations
  # under the License.
  """Pipeline executor that executes a series of modules in a pipeline fashion."""
  import json
  import os
  import tvm._ffi
  from tvm import relay
  from tvm.relay.transform import InferType
  from tvm.contrib.pipeline_executor import PipelineExecutorFactoryModule
  
  
  def pipeline_executor_build_enabled():
      """Check if the pipeline executor build is enabled.
  
      Return
      -------
      enable: bool
          Return whether the pipeline executor is enabled.
      """
      return tvm.contrib.pipeline_executor.pipeline_executor_enabled()
  
  
  def build(pipe_configs):
      """Build modules used in the pipeline executor, then use these modules and configuration
      to create a pipeline executor.
  
      Parameters
      ----------
      pipe_configs: PipelineConfig
          Build Configuration information.
  
      Returns
      -------
      ret: PipelineExecutorFactoryModule
          Common interface for pipeline executor factory modules.
      """
      libs = {}
      config = pipe_configs.get_config()
      if "module_connection" not in config:
          raise RuntimeError('"module_connection" is missing')
      if "input_connection" not in config:
          raise RuntimeError('"input_connection" is missing')
      if "param_connection" not in config:
          raise RuntimeError('"param_connection" is missing')
  
      mod_n_configs = config["module_connection"]
      config_len = len(mod_n_configs)
      module_string_config = [{} for _ in range(config_len)]
      # Use hardware configurations to build backend modules for each subgraph.
      for ir_mod, mod_config in mod_n_configs.items():
          pipe_config = mod_config["pipeline"].copy()
          mod_idx = pipe_config["mod_idx"]
          dev = mod_config["dev"]
          target = mod_config["target"]
          build_func = relay.build
          # Callers may need to use a customized building function to wrap the pre-building logic
          # and the backend building logic. For example, in order to support a backend which only
          # can do "int8" computation, the caller may need to merge the "quantization" logic
          # into the building logic to creat a customized building function.
          if "build" in mod_config and mod_config["build"]:
              build_func = mod_config["build"]
  
          lib = build_func(
              ir_mod,
              target,
              params=mod_config["params"],
              target_host=mod_config["target_host"],
              mod_name=mod_config["mod_name"],
          )
  
          pipe_config["dev"] = "{},{}".format(dev.device_type, dev.device_id)
          # Use "mod_idx" as the key to create a "module_connection" map which is not only
          # for the module index but also for the module connection used to build the pipeline.
          module_string_config[mod_idx] = pipe_config
          libs[mod_idx] = {"lib": lib, "dev": dev, "fcompile": mod_config["fcompile"]}
  
      # Creating a text form configuration to record the "input_connection" and the
      # "module_connection" information. The "input_connection" is used to record the
      # map of global input and subgraph input, and the "module_connection" is used to
      # record module dependency.
      string_config = {}
      string_config["param_connection"] = config["param_connection"]
      string_config["input_connection"] = config["input_connection"]
      string_config["module_connection"] = module_string_config
  
      return PipelineExecutorFactoryModule(libs, string_config)
  
  
  def export_library(factory, directory_path):
      """Export the pipeline executor into disk files.
  
      Parameters
      ----------
      factory : PipelineExecutorFactoryModule
          The pipeline executor factory
      directory_path : str
          Export the files to this directory.
      """
      if not factory.pipeline_mods:
          raise RuntimeError("The pipeline executor has not been initialized.")
  
      # Check if the directory_path exists.
      if not directory_path or not os.path.exists(directory_path):
          raise RuntimeError("The directory {directory_path} does not exist.")
      # Create an load configuration.
      load_config_file_name = "{}/load_config".format(directory_path)
      pipeline_config_file_name = "{}/pipeline_config".format(directory_path)
      config = {}
      config["load_config"] = load_config_file_name
      config["pipeline_config"] = pipeline_config_file_name
      load_config = []
      # Export the library, JSON, and parameter into files, then export these files path
      # into a configuration file.
      for lib_index in factory.pipeline_mods:
          mconfig = {}
          mconfig["mod_idx"] = lib_index
          mconfig["lib_name"] = "{}/lib{}.so".format(directory_path, lib_index)
          mconfig["json_name"] = "{}/json{}".format(directory_path, lib_index)
          mconfig["params_name"] = "{}/params{}".format(directory_path, lib_index)
          lib_config = factory.pipeline_mods[lib_index]
          mconfig["dev"] = "{},{}".format(
              lib_config["dev"].device_type,
              lib_config["dev"].device_id,
          )
          fcompile = lib_config["fcompile"]
          if not fcompile:
              fcompile = False
  
          # Get the graph, lib, and parameters from GraphExecutorFactoryModule.
          lib = factory.pipeline_mods[lib_index]["lib"]
          # Export the lib, graph, and parameters to disk.
          lib.export_library(mconfig["lib_name"], fcompile)
          with open(mconfig["json_name"], "w") as file_handle:
              file_handle.write(lib.graph_json)
          with open(mconfig["params_name"], "wb") as file_handle:
              file_handle.write(relay.save_param_dict(lib.params))
  
          load_config.append(mconfig)
  
      with open(load_config_file_name, "w") as file_handle:
          json.dump(load_config, file_handle)
  
      with open(pipeline_config_file_name, "w") as file_handle:
          json.dump(factory.mods_config, file_handle)
  
      config_file_name = "{}/config".format(directory_path)
      with open(config_file_name, "w") as file_handle:
          json.dump(config, file_handle)
  
      return config_file_name
  
  
  class PipelineConfig(object):
      """Pipeline configuration information, this class contains the DAG that expresses
      the dependency of each module involved in a pipeline and the parameters for building
      each module.
      """
  
      class Binding:
          """This class defines the module connections information.
          The type can only be "input" or "output".
  
          Parameters
          ----------
          owner : ModuleWrapper
              The class who owns this interface.
  
          io_type : str
              The I/O type of this interface. It can only be "input" or "output".
  
          name : str/integer
              Name, for input it is string such as "data0", for output it is an integer such as 0.
  
          data_type: TensorType
              The data type of this interface.
          """
  
          def __init__(self, owner, io_type, name, data_type=None):
              self.io_owner = owner
              self.io_type = io_type
              self.name = str(name)
              # Child interfaces that depend on this interface.
              self.bindings = []
              # Parents interfaces that this interface depend on.
              self.parents = []
  
              self.data_type = data_type
  
          def get_name(self):
              # Return name of this interface and the name of owner who owns this interface.
              owner_name = ""
              if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                  owner_name = self.io_owner.name
  
              return owner_name, self.name
  
          def get_owner_idx(self):
              # If the owner is ModuleWrapper return the owner index, if not return 0.
              if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                  return self.io_owner.idx
  
              return -1
  
          def is_pipeline_executor_interface(self):
              """The pipeline interface is used to interact with the caller. There are two types
              of interfaces, one is 'input' another is 'output'. The pipeline input interface
              is responsible for passing parameters to the internal module interface, and the
              pipeline output interface is responsible for outputting the results computed by
              the pipeline executor to the caller.
              """
              return not isinstance(self.io_owner, PipelineConfig.ModuleWrapper)
  
          def __repr__(self):
              # Geting the binding information in the form of text.
              str_format = "  |{}: ".format(self.name)
              for binding in self.bindings:
                  mname, dname = binding.get_name()
                  str_format += "{0}:{1} ".format(mname, dname)
  
              return str_format
  
          def check_binding_dict(self, connection_dict):
              """Checking the binding dictionary.
              Parameter
              ---------
              connection_dict : Dict[str, Any]
                  It is a dictionary of module connections.
              """
              if "interface_name" not in connection_dict:
                  raise RuntimeError('"inteface_name" is missing in global config!"')
              if "connection" not in connection_dict:
                  raise RuntimeError(f'"connection" is missing!"')
              # The global interface mapping should be one-to-one.
              if not connection_dict["connection"]:
                  raise RuntimeError("The global interface map is empty!")
              if len(connection_dict["connection"]) > 1:
                  raise RuntimeError("A global interface maps multiple module interfaces!")
              if "mod_idx" not in connection_dict["connection"][0]:
                  raise RuntimeError('"mod_idx" is missing!')
  
          def get_binding_dict(self):
              """Returning the binding information in the form of dictionary.
              Returns
              -------
              data : Dict[str, Any]
                  The binding information is in the form of dictionary.
              """
              dict_format = {"interface_name": self.name, "connection": []}
              for binding in self.bindings:
                  _, dname = binding.get_name()
                  midx = binding.get_owner_idx()
                  dict_format["connection"].append({"mod_idx": midx, "interface_name": dname})
  
              self.check_binding_dict(dict_format)
              return dict_format
  
          def check_dag_acyclic(self, start, inputs):
              """This is to check whether the DAG containing these input interfaces is acyclic.
              Parameters
              ----------
              start: ModuleWrapper
                  The starting node of the cycle check algorithm.
  
              inputs: Binding
                  These interfaces are used to connect to each other to build DAG.
  
              Return
              ------
                  Return true if there is no cycle in the DAG.
              """
              for binding in inputs.values():
                  if start == binding.io_owner:
                      return False
                  for p in binding.parents:
                      if not self.check_dag_acyclic(start, p.io_owner.input_bindings.bindings):
                          return False
  
              return True
  
          def connect(self, binding):
              """Connect the current interface to the destination interface.
              Correct connections are as follows: 1. the pipeline input connected to a module input,
              2. the module output connected to a pipeline output, 3. the module output connected to
              a module input.
  
              Parameters
              ----------
              binding: Binding
                  The destination of this connection.
              """
  
              # Check whether the binding setting is correct or not.
              if self.io_owner == binding.io_owner:
                  raise RuntimeError("Can not bind itself.")
  
              if self.io_type == "param" and not self.is_pipeline_executor_interface():
                  raise RuntimeError(
                      'The "param" binding can only be used by a pipeline executor interface!'
                  )
  
              if not self.is_pipeline_executor_interface() and self.io_type == "input":
                  raise RuntimeError("Module can only bind from output interface!")
  
              if self.io_type == "param" and binding.io_type != "param":
                  raise RuntimeError(
                      'A global "param" interface can only be bind with a module "param" interface!'
                  )
  
              if (
                  not self.is_pipeline_executor_interface()
                  and not binding.is_pipeline_executor_interface()
                  and binding.io_type == "output"
              ):
                  raise RuntimeError("Can not bind module output with another module output!")
  
              if (
                  not self.is_pipeline_executor_interface()
                  and binding.is_pipeline_executor_interface()
                  and binding.io_type == "input"
              ):
                  raise RuntimeError("Can not bind module output with pipeline input!")
  
              if self.is_pipeline_executor_interface() and self.io_type == "output":
                  raise RuntimeError("Global output can not be used as binding start point.")
  
              if (
                  self.is_pipeline_executor_interface()
                  and self.io_type == "input"
                  and binding.io_type != "input"
              ):
                  raise RuntimeError("Global input can only bind with module input.")
  
              self.bindings.append(binding)
              if not self.is_pipeline_executor_interface():
                  # Check whether the data types of the source and destination are the same.
                  if (
                      isinstance(binding.io_owner, PipelineConfig.ModuleWrapper)
                      and self.data_type != binding.data_type
                  ):
                      raise RuntimeError(
                          f"Illegal type (%s vs. %s): binding type is not same!"
                          % (self.data_type, binding.data_type)
                      )
  
                  binding.parents.append(self)
  
                  # Do acyclic check after increasing the in-degree of child node by setting
                  # current interface as a parent of the child node.
  
                  if not self.check_dag_acyclic(
                      binding.io_owner, self.io_owner.input_bindings.bindings
                  ):
                      raise RuntimeError("Illegal connection: Cause a cycle!")
  
      class BindingList:
          """Container for bindings(input or output interface).
  
          Parameters
          ----------
          owner : ModuleWrapper/PipelineConfig
              The owner of this class can be ModuleWrapper or PipelineConfig.
  
          io_type : str
              The type of this class can be "input" or "output".
          """
  
          def __init__(self, owner, io_type):
              self.bindings = {}
              self.io_owner = owner
              self.binding_type = io_type
  
          def get_binding_data_type(self, key):
              if isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                  return self.io_owner.get_data_type(key, self.binding_type)
              return None
  
          def __getitem__(self, key):
              if key not in self.bindings:
                  data_type = self.get_binding_data_type(key)
                  if not data_type and isinstance(self.io_owner, PipelineConfig.ModuleWrapper):
                      raise RuntimeError(f"Can not find {key} in binding list {self.binding_type}.")
  
                  self.bindings[key] = PipelineConfig.Binding(
                      self.io_owner, self.binding_type, key, data_type
                  )
  
              return self.bindings[key]
  
      class ModuleWrapper:
          """This class is a wrapper representing the module and contains information such as
          module information, binding information and building information.
          """
  
          def __init__(self, mod=None):
              self.target_host = None
              self.build_func = None
              self.params = None
              self.target = None
              self.fcompile = None
              self.name = None
              self.dev = None
              self.cpu_affinity = ""
              self.idx = None
              self.mod = mod
              self.input_params = InferType()(mod)["main"].params
              self.output_type = InferType()(mod)["main"].checked_type.ret_type
              self.input_bindings = PipelineConfig.BindingList(self, "input")
              self.output_bindings = PipelineConfig.BindingList(self, "output")
              self.param_binding = PipelineConfig.Binding(self, "param", "param")
  
          def __eq__(self, other):
              if isinstance(other, PipelineConfig.ModuleWrapper):
                  return self.mod == other.mod
  
              return False
  
          def __getitem__(self, key):
              if isinstance(key, str):
                  if key == "input":
                      return self.input_bindings
  
                  if key == "output":
                      return self.output_bindings
  
                  if key == "param":
                      return self.param_binding
  
                  raise RuntimeError(f"{key} not found!")
  
              raise RuntimeError('The data type of "key" is not supported!')
  
          def get_data_type(self, key, interface_type):
              """Get the module interface data type according to the key value and interface type.
              Parameters
              ----------
              key: str
                  The interface name.
  
              interface_type:
                  The interface type.
  
              Return
              -------
                  Return data type.
              """
              if interface_type == "input":
                  for param in self.input_params:
                      if param.name_hint == key:
                          return param._checked_type_
  
              if interface_type == "output":
                  if isinstance(self.output_type, tvm.ir.type.TupleType):
                      if int(key) < len(self.output_type.fields):
                          return self.output_type.fields[int(key)]
                  elif int(key) == 0:
                      return self.output_type
  
              return None
  
          def set_idx_name(self, idx):
              # Set the index value and generate the module name.
              self.idx = idx
              self.name = "mod{}".format(str(idx))
  
          def is_root_mod(self):
              """Check whether this node is the root node in DAG, this function is used
              in topological sort.
              """
              return all([not b.parents for b in self.input_bindings.bindings.values()])
  
          def remove_self_from_bindings(self):
              """Remove the current node from child dependencies to reduce the in-degree
              of child node, this function is used in topological sort.
              """
              for binding in self.output_bindings.bindings.values():
                  for child in binding.bindings:
                      if binding in child.parents:
                          child.parents.remove(binding)
  
      def __init__(self):
          self.mod_wrapper = {}
          self.input_bindings = self.BindingList(self, "input")
          self.output_bindings = self.BindingList(self, "output")
          # There is a map of global parameters group and module index.
          self.param_group_bindings = self.BindingList(self, "param")
  
      def __str__(self):
          # Get configuration information as a string.
  
          # Use topological sort to get correct module order.
          self.dag_topology_sort()
          # Getting the parameters dependencies.
          param_dump = "Params\n"
          for param_name in self.param_group_bindings.bindings:
              inf = self.param_group_bindings.bindings[param_name]
              param_dump += str(inf) + "\n"
          # Get the input dependencies.
          input_dump = "\nInputs\n"
          for input_name in self.input_bindings.bindings:
              inf = self.input_bindings.bindings[input_name]
              input_dump += str(inf) + "\n"
  
          # Get the connections information of each module.
          output = {}
          connections_dump = "\nconnections\n"
          for mod in self.mod_wrapper:
              for interface in self.mod_wrapper[mod].output_bindings.bindings.values():
                  if interface.bindings:
                      mname, dname = interface.get_name()
                      iname = mname + ".output(" + dname + ")->"
                      for dep in interface.bindings:
                          dep_mname, dep_dname = dep.get_name()
                          if isinstance(dep.io_owner, PipelineConfig.ModuleWrapper):
                              iname += f" {dep_mname}.{dep_dname}"
                              connections_dump += f"  |{iname}\n"
                          else:
                              output[dep_dname] = f"{mname}.output({dname})"
  
          # Get the output dependencies.
          output_dump = "\noutput\n"
          for name in sorted(output.keys()):
              output_dump += f"  |output({name}) : {output[name]}\n"
  
          return param_dump + input_dump + output_dump + connections_dump
  
      def __getitem__(self, key):
          if isinstance(key, tvm.ir.module.IRModule):
              if key not in self.mod_wrapper:
                  self.mod_wrapper[key] = self.ModuleWrapper(key)
              return self.mod_wrapper[key]
  
          if isinstance(key, str):
              if key == "input":
                  return self.input_bindings
              if key == "output":
                  return self.output_bindings
              if key == "param_group":
                  return self.param_group_bindings
  
              raise RuntimeError(f"{key} not found!")
  
          raise RuntimeError(f'The key type "{type(key)}" is not supported!')
  
      def get_config(self):
          """Get the configuration information in dictionary form, this configuration
          will be used to create pipeline executor.
          """
  
          # Use topological sort to get the correct order of modules.
          self.dag_topology_sort()
          mconfig = {}
          module_connection = {}
          for mod in self.mod_wrapper:
              # Generate pipeline configuration.
              mconf = {}
              output_conf = []
              module = self.mod_wrapper[mod]
              for _, binding in module.output_bindings.bindings.items():
                  dep_conf = []
                  output = {}
                  if binding.bindings:
                      for dep in binding.bindings:
                          dep_item = {}
                          _, dname = dep.get_name()
                          if dep.is_pipeline_executor_interface():
                              dep_item["global_output_index"] = int(dname)
                          else:
                              dep_item["mod_idx"] = dep.get_owner_idx()
                              dep_item["input_name"] = dname
                          dep_conf.append(dep_item)
  
                  # The value of ouput_idx start from 0.
                  output["output_idx"] = int(binding.name)
                  output["dependencies"] = dep_conf
                  output_conf.append(output)
  
              mconf["mod_idx"] = module.idx
              mconf["cpu_affinity"] = module.cpu_affinity
              mconf["output"] = output_conf
  
              module_connection[mod] = {
                  "pipeline": mconf,
                  "target_host": module.target_host,
                  "mod_name": "default",
                  "build": module.build_func,
                  "params": module.params,
                  "target": module.target,
                  "fcompile": module.fcompile,
                  "dev": module.dev,
              }
  
          # Creating a map including pipeline inputs and subgraph inputs.
          input_connection = []
          for input_name in self.input_bindings.bindings:
              input_dict = self.input_bindings.bindings[input_name].get_binding_dict()
              if "interface_name" not in input_dict["connection"][0]:
                  raise RuntimeError("interface_name is missing in connection config!")
              # Creating the map including global interfaces and subgraph interfaces.
              input_map = {
                  "global_interface_name": input_dict["interface_name"],
                  "mod_idx": input_dict["connection"][0]["mod_idx"],
                  "module_interface_name": input_dict["connection"][0]["interface_name"],
              }
              input_connection.append(input_map)
  
          # Create a map including global parameters groups and modules.
          param_connection = []
          for param_name in self.param_group_bindings.bindings:
              param_dict = self.param_group_bindings.bindings[param_name].get_binding_dict()
              param_map = {
                  "global_param_name": param_dict["interface_name"],
                  "mod_idx": param_dict["connection"][0]["mod_idx"],
              }
              param_connection.append(param_map)
  
          mconfig["module_connection"] = module_connection
          mconfig["input_connection"] = input_connection
          mconfig["param_connection"] = param_connection
          return mconfig
  
      def dag_topology_sort(self):
          """Use topological sort to get order of pipeline modules."""
          mlist = []
          mod_wrapper = self.mod_wrapper.copy()
          while mod_wrapper:
              temp_list = []
              for mod, wrapper in mod_wrapper.items():
                  if wrapper.is_root_mod():
                      temp_list.append(mod)
                      wrapper.remove_self_from_bindings()
  
              for mod in temp_list:
                  mod_wrapper.pop(mod, None)
  
              mlist += temp_list
  
          mod_wrapper_sort = {}
          for mod, i in zip(mlist, range(len(mlist))):
              self.mod_wrapper[mod].set_idx_name(i)
              mod_wrapper_sort[mod] = self.mod_wrapper[mod]
  
          self.mod_wrapper = mod_wrapper_sort
  
      def get_mod_idx(self, mod):
          # Return the module index.
          idx = self.mod_wrapper[mod].idx
          return idx
  
      def pipe_input(self, name):
          # Return the input interface according to the name.
          return self.input_bindings[name]
  
      def pipe_output(self, idx):
          # Return the output interface according to the name.
          return self.output_bindings[idx]
  