#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/host_context/async_value.h"  // from @tf_runtime
#include "tfrt/host_context/chain.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime

        [
            "//third_party/tf_runtime_llvm:tfrt_translate_main",
        ],
        ["@tf_runtime//third_party/llvm_derived:tfrt_translate_main"],
    ),
)

tf_cc_binary(
    name = "bef_executor",
    testonly = True,
    visibility = [":friends"],
    deps = [
        "@tf_runtime//:dtype",
        "@tf_runtime//:simple_tracing_sink",
        "@tf_runtime//tools:bef_executor_expensive_kernels",

    return isa<tf_executor::IslandOp>(region->getParentOp());
  }
};

}  // namespace

TensorFlowExecutorDialect::TensorFlowExecutorDialect(MLIRContext *context)
    : Dialect(/*name=*/"tf_executor", context,
              TypeID::get<TensorFlowExecutorDialect>()) {
  addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc.inc"
      >();

func.func @simple_independent_chains(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>) {
  // CHECK: tf_executor.graph {
  tf_executor.graph {
    // CHECK: %[[A_CONTROL:.*]] = tf_executor.island wraps "tf.OpA"() : () -> ()
    %control_A = tf_executor.island wraps "tf.OpA"() : () -> ()
    // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>

    %control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()
    %ri, %c0 = tf_executor.island wraps "tf.TPUReplicatedInput"(%arg0, %arg1) {index = 1 : i64, is_mirrored_variable = false, is_packed = false} : (tensor<i32>, tensor<i32>) -> tensor<i32>

// CHECK: TPUCompile
// CHECK: TPUExecute
// CHECK-NOT: func @_func
    tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "arith.constant"() {value = dense<2.000000e+00> : tensor<f32>} : () -> tensor<f32>
      %outputs_0, %control_1 = tf_executor.island wraps "arith.constant"() {value = dense<3.000000e+00> : tensor<f32>} : () -> tensor<f32>

    %0 = tf_executor.graph {
      %1, %c0 = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor<i1>} : () -> tensor<i1>
      tf_executor.fetch %1 : tensor<i1>
    }
    return %0 : tensor<i1>
  }

  func.func private @while_body(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
    tf_executor.graph {

  // CHECK-NOT: _replication_info =
  %0 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.GuaranteeConst"(%arg1) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %outputs_0, %control_1 = tf_executor.island wraps "tf.GuaranteeConst"(%arg2) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %control_2 = tf_executor.island wraps "tf.NoOp"() {_pivot_for_cluster = "cluster", device = ""} : () -> ()

    auto executor_graph =
        rewriter.create<tf_executor::GraphOp>(loc, func_type.getResults());
    rewriter.inlineRegionBefore(graph.getNodes(), executor_graph.getBody(),
                                executor_graph.getBody().end());

    // Add terminator of tf_executor::graph
    rewriter.setInsertionPointToEnd(&executor_graph.getBody().front());
    rewriter.create<tf_executor::FetchOp>(loc);

    // Add terminator of func

    tf_executor.graph {
      %control = tf_executor.island {
        "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @partitioned}> : (tensor<*x!tf_type.resource>) -> ()
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }