#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_remaining_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/verify_suitable_for_graph_export.h"

namespace mlir {
namespace tf_executor {
namespace {

// Comparator for `OpsInReverseProgramOrder`.

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

#define DEBUG_TYPE "tf-executor-tpu-v1-island-inlining"

namespace mlir {
namespace tf_executor {

namespace {

#         "//tensorflow/core/runtime_fallback/util:tensor_util",
#         "//tensorflow/core/tfrt/utils:fallback_tensor",
#         "@eigen_archive//:eigen3",
#         "@tf_runtime//:bef",
#         "@tf_runtime//:befexecutor",
#         "@tf_runtime//:core_runtime_alwayslink",
#         "@tf_runtime//:hostcontext_alwayslink",
#         "@tf_runtime//:mlirtobef",
#         "@tf_runtime//:support",
#         "@tf_runtime//:tensor",

# CHECK: %[[output:.*]], %[[control:.*]]tf_executor.island wraps "tf.ParseExample"
# CHECK: operandSegmentSizes = array<i32: 1, 1, 0, 1, 1>
# CHECK: resultSegmentSizes = array<i32: 0, 0, 0, 1>
# CHECK: tf_executor.fetch %[[output]] : tensor<*xi64>

node {
  name: "serilaized"
  op: "Const"
  attr {
    key: "dtype"
    value {

#include "tensorflow/compiler/mlir/init_mlir.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

int main(int argc, char **argv) {
  tensorflow::InitMlir y(&argc, &argv);

// Creates a pass that converts mlir functions consisting of mlir ops into a
// tf_executor dialect as a single island.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateFunctionalToExecutorDialectConversionPass();

// Creates a pass that lifts inner ops of tf_executor.island ops in
// tf_executor.graph into the same block as the tf_executor.graph.
std::unique_ptr<OperationPass<func::FuncOp>>

  %graph:2 = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<1.0> : tensor<1x2xf16>} : () -> tensor<1x2xf16> loc("const1")
    %1:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<[1.0, 2.0]> : tensor<2xf16>} : () -> tensor<2xf16> loc("const2")
    tf_executor.fetch %0#0, %1#0 : tensor<1x2xf16>, tensor<2xf16>
  }

// Lowers MLIR module to XLA HLO inside an XlaComputation. The input module
// should only contain operations in tf dialect. If the input module contains
// operation in the tf_executor dialect, for example, returns an error.
// Exception to this are tf_executor dialect ops that are optimized away through
// canonicalization.
//
// Operations in tf dialect are lowered to XLA HLO through the following steps:

// outputs to generate a single control output.
//
// For example,
// ```
// %merged_control = "tf_executor.island"(%control_a, %control_b) ({
//   "tf.NoOp"() : () -> ()
//   "tf_executor.yield"() : () -> ()
// }) : (!tf_executor.control, !tf_executor.control) -> (!tf_executor.control)
// ```
//
// `%merged_control` is a NoOp control barrier in this case.
//

        new llvm::SmallDenseSet<mlir::TypeID, 32>{
            TypeID::get<mlir::ModuleOp>(),
            TypeID::get<mlir::tf_executor::GraphOp>(),
            TypeID::get<mlir::func::ReturnOp>(),
            TypeID::get<mlir::func::FuncOp>(),
            TypeID::get<mlir::tf_executor::YieldOp>(),
            TypeID::get<mlir::tf_executor::IslandOp>(),
            TypeID::get<TF::TPUReplicatedInputOp>(),
            TypeID::get<TF::TPUReplicatedOutputOp>(),