// 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.

// Test for partial folding: only fold i32 constants.
// RUN: tac-opt-all-backends -tfl-fold-constants-to-subgraph='fold-all-constants=false' %s -split-input-file -verify-diagnostics | FileCheck --check-prefix=PARTIAL %s

// Test for fold all constants.
// RUN: tac-opt-all-backends -tfl-fold-constants-to-subgraph='fold-all-constants=true' %s -split-input-file -verify-diagnostics | FileCheck --check-prefix=ALL %s

module {

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

tf_cc_binary(
    name = "tac-opt-all-backends",
    testonly = True,
    deps = [
        ":tac-opt_lib",
        "//tensorflow/compiler/mlir/lite/experimental/common:outline_operations",
        "//tensorflow/compiler/mlir/lite/experimental/tac/hardwares:all-target-hardwares",
    ],
)

cc_library(
    name = "tac-translate-lib",
    srcs = [
        "tac_translate.cc",
    ],
    deps = [

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

#include "tensorflow/compiler/mlir/lite/experimental/tac/common/cost.h"
#include "tensorflow/compiler/mlir/lite/experimental/tac/common/subgraph.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/transforms/cost_model.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/utils/utils.h"

namespace mlir {
namespace TFL {

}

// CHECK: %[[CST:.*]] = arith.constant dense<1> : tensor<4xi32>
// CHECK:  [[VAL_0:%.*]] = "tfl.reshape"(%1, %[[CST]]) {tac.device = "GPU",  tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:  [[VAL_1:%.*]] = "tfl.reshape"(%2, %[[CST]]) {tac.device = "GPU",  tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>

//
// ======== IRs after TAC ========
// op1 = .... {tac.device = "CPU"}
// op2 = .... {tac.device = "GPU"}
// op3 = .... {tac.device = "GPU"}
// op4 = .... {tac.device = "CPU"}
//
// We will run a separate cost estimation for each op with each different
// hardwares.
//
// ======== IRs after cost-estimation ========
// op1 = .... {tac.device = "CPU"} {"CPU": 10, "GPU": -1.0}