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.
  # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name, W0611, C0413
  
  """Expose Vitis-AI test functions to the Python frontend"""
  
  import sys
  import numpy as np
  
  import pytest
  
  pytest.importorskip("pyxir")
  import pyxir.contrib.target.DPUCADX8G
  import pyxir.contrib.target.DPUCZDX8G
  
  import tvm
  from tvm import relay
  from tvm import runtime
  from tvm.relay import transform
  from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
  from tvm.relay.build_module import bind_params_by_name
  from tvm.contrib.target import vitis_ai
  from tvm.contrib import graph_executor
  from tvm.contrib import utils
  
  
  def get_cpu_op_count(mod):
      """Traverse graph counting ops offloaded to TVM."""
  
      class Counter(tvm.relay.ExprVisitor):
          def __init__(self):
              super().__init__()
              self.count = 0
  
          def visit_call(self, call):
              if isinstance(call.op, tvm.ir.Op):
                  self.count += 1
  
              super().visit_call(call)
  
      c = Counter()
      c.visit(mod["main"])
      return c.count
  
  
  def skip_test():
      """Skip test if it requires the Vitis-AI codegen and it's not present."""
      if not tvm.get_global_func("relay.ext.vitis_ai", True):
          print("Skip test because Vitis-AI codegen is not available.")
          return True
      return False
  
  
  def build_module(
      mod,
      target,
      dpu_target="DPUCADF8H",
      params=None,
      enable_vitis_ai=True,
      tvm_ops=0,
      vitis_ai_partitions=1,
  ):
      """Build module for Vitis-AI codegen."""
      if isinstance(mod, tvm.relay.expr.Call):
          mod = tvm.IRModule.from_expr(mod)
      if params is None:
          params = {}
  
      with tvm.transform.PassContext(
          opt_level=3, config={"relay.ext.vitis_ai.options.target": dpu_target}
      ):
          if enable_vitis_ai:
              mod = partition_for_vitis_ai(mod, params, dpu_target)
              tvm_op_count = get_cpu_op_count(mod)
              assert tvm_op_count == tvm_ops, "Got {} TVM operators, expected {}".format(
                  tvm_op_count, tvm_ops
              )
              partition_count = 0
              for global_var in mod.get_global_vars():
                  if "vitis_ai" in global_var.name_hint:
                      partition_count += 1
  
              assert (
                  vitis_ai_partitions == partition_count
              ), "Got {} Vitis-AI partitions, expected {}".format(
                  partition_count, vitis_ai_partitions
              )
          relay.backend.te_compiler.get().clear()
          return relay.build(mod, target, params=params)
  
  
  def update_lib(lib, cross_compile=None):
      tmp_path = utils.tempdir()
      lib_name = "lib.so"
      lib_path = tmp_path.relpath(lib_name)
      if cross_compile:
          lib.export_library(lib_path, cc=cross_compile)
      else:
          lib.export_library(lib_path)
      lib = runtime.load_module(lib_path)
      return lib
  
  
  def extract_vitis_ai_modules(module):
      """Get the Vits-AI runtime module from llvm module."""
      return list(
          filter(lambda mod: mod.type_key == "VitisAIRuntime", module.get_lib().imported_modules)
      )
  
  
  def verify_codegen(
      module, num_vitis_ai_modules=1, params=None, target="llvm", tvm_ops=0, dpu_target="DPUCADX8G"
  ):
      """Check Vitis-AI codegen against a known good output."""
      module = build_module(
          module,
          target,
          params=params,
          dpu_target=dpu_target,
          tvm_ops=tvm_ops,
          vitis_ai_partitions=num_vitis_ai_modules,
      )
      vitis_ai_modules = extract_vitis_ai_modules(module)
  
      assert len(vitis_ai_modules) == num_vitis_ai_modules, (
          f"The number of Vitis-AI modules produced ({len(vitis_ai_modules)}) does not "
          f"match the expected value ({num_vitis_ai_modules})."
      )
  
  
  def verify_result(
      mod,
      map_inputs,
      out_shape,
      result,
      tol=1e-5,
      target="llvm",
      device=tvm.cpu(),
      params=None,
      dpu_target="DPUCADX8G",
      tvm_ops=0,
  ):
      """To check the result between reference and byoc vitis-ai flow"""
  
      lib = build_module(mod, target, params=params, dpu_target=dpu_target, tvm_ops=tvm_ops)
      lib = update_lib(lib)
      rt_mod = graph_executor.GraphModule(lib["default"](tvm.cpu()))
  
      for name, data in map_inputs.items():
          rt_mod.set_input(name, data)
      rt_mod.set_input(**params)
      rt_mod.run()
  
      out_shapes = out_shape if isinstance(out_shape, list) else [out_shape]
      results = result if isinstance(result, list) else [result]
  
      for idx, shape in enumerate(out_shapes):
          out = tvm.nd.empty(shape, device=device)
          out = rt_mod.get_output(idx, out)
          tvm.testing.assert_allclose(out.numpy(), results[idx], rtol=tol, atol=tol)
  