#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/platform/path.h"

namespace mlir {
namespace tf_saved_model {
namespace {

#define GEN_PASS_DEF_FREEZEASSETSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_savedmodel_passes.h.inc"

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_executor::FetchOp;
using ::mlir::tf_executor::GraphOp;
using ::mlir::tf_executor::IslandOp;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::tensorflow::kImportModelDefaultGraphFuncName;

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {

std::vector<llvm::StringRef> GetEntryFunctionAttributeNames() {
  return {"tf.entry_function",
          tf_saved_model::kTfSavedModelInitializerTypeAttr};
}

bool IsEntryFunction(func::FuncOp func) {
  for (const auto &attr : GetEntryFunctionAttributeNames()) {

  # has a corresponding tf.LegacyCall op.
  # CHECK:      func {{@[a-zA-Z_0-9]+}}(
  # CHECK-SAME: [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["y"]},
  # CHECK-SAME: [[ARG1:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]}
  #
  # CHECK-NEXT: [[R0:%.*]] = "tf.LegacyCall"([[ARG1]], [[ARG0]])
  z = plus(x, y)

  tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)

  func.func @main(%arg0: tensor<10x1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<10x1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/core/public/session.h"

namespace mlir {
namespace tf_saved_model {

// Freezes readonly variables in the graph.
LogicalResult FreezeVariables(ModuleOp module, tensorflow::Session* session);

}  // namespace tf_saved_model

}  // namespace mlir

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

// A pass that removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns

#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace mlir {
namespace tf_saved_model {

// Create a pass that removes function arguments that map to global tensors.
std::unique_ptr<Pass> CreateLowerGlobalsToMlProgramPass();

// Register this pass in the global registry of MLIR.
void RegisterLowerGlobalsToMlProgramPass();

}  // namespace tf_saved_model
}  // namespace mlir

// RUN: tf-tfrt-opt -verify-diagnostics -split-input-file -tfrt-fuse-tpu-compile-and-execute-ops -canonicalize %s | FileCheck %s --dump-input=fail --dump-input-filter=all

module attributes {tf_saved_model.semantics} {

// Test fusing _TPUCompileMlirOp and TPUExecuteOp into TPUCompileMlirAndExecuteOp.

// CHECK-LABEL: func private @test_fuse_tpu_ops
func.func private @test_fuse_tpu_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {

  func.func @conv_with_two_layers(%arg0: tensor<1x3x4x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<1x3x2x2xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {