# limitations under the License.
# ==============================================================================
"""Target aware conversion for TFLite model."""

from tensorflow.compiler.mlir.lite.experimental.tac.py_wrapper import _pywrap_tac_wrapper


def run_tac(model_path, targets, output_path):
  """Run target aware conversion for the given tflite model file.

  Args:
    model_path: Path to the tflite model file.

// RUN: tac-opt-all-backends -tfl-tac-filter='use-test-setting=true' %s -split-input-file -verify-diagnostics | FileCheck %s

// expected-remark@below {{Tac filter (0): filter type: function filter SKIP_TARGET_ANNOTATION, filter_pattern: "^testFunction"}}
// expected-remark@below {{Tac filter (1): filter type: function filter INCLUDE_TARGET_ANNOTATION, filter_pattern: "testFunctionInclude"}}
// expected-remark@below {{Tac filter (1) specified but not applied to any op}}

  // CHECK: tac.cost = 1.000000e+03
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  // CHECK: tac.cost = 1.000000e+03
  %1 = "tfl.mul"(%0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>

   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %1 = "tfl.add"(%arg0, %0) {fused_activation_function = "RELU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
   // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"

#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_importer_exporter.h"

namespace mlir {
namespace TFL {
namespace tac {

// Main class for using Target Aware Conversion (TAC).
// To run TAC:
// 1) users should create object form this class, with desired options
//    (TacModule::Options).

  pass_manager->addPass(mlir::TFL::tac::CreateTargetAnnotationPass(this));
  pass_manager->addPass(mlir::TFL::tac::CreateRaiseTargetSubgraphsPass());
  pass_manager->addPass(mlir::TFL::tac::CreateFoldConstantsToSubgraphPass(
      /*fold_all_constants=*/false));
  pass_manager->addPass(
      mlir::TFL::tac::CreateAlternativeSubgraphPass(device_specs));
  if (options_.legalize_to_tflite_ops) {

// RUN: tac-opt-all-backends -tfl-compute-cost %s -split-input-file -verify-diagnostics | FileCheck %s

// CHECK: tac.cost = 7.864320e+05
func.func @func_0_CPU(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32> attributes {tac.device = "CPU", tac.interface_name = "func_0"} {
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<256x32x32x3xf32>, tensor<256x32x32x3xf32>) -> tensor<256x32x32x3xf32>

#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"

namespace mlir {
namespace TFL {
namespace tac {
// An OperationPass<> with access to the TAC module instance that the
// pass is running part of.

#include "tensorflow/compiler/mlir/lite/experimental/tac/py_wrapper/tac_wrapper.h"

#include <memory>
#include <string>
#include <vector>

#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"

#include "tensorflow/compiler/mlir/lite/experimental/tac/common/targets.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/utils.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/execution_metadata_exporter.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/target_hardware.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/tac_module.h"