#include "llvm/ADT/StringRef.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/threadpool_interface.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
#include "tfrt/host_context/host_context.h"  // from @tf_runtime

#include "tensorflow/core/runtime_fallback/util/fallback_test_util.h"
#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
#include "tfrt/core_runtime/core_runtime.h"  // from @tf_runtime
#include "tfrt/host_context/async_value.h"  // from @tf_runtime
#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime

#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

 public:
  using Base::Base;
  static constexpr ::mlir::StringLiteral name = "tf_executor.control";
};

class TokenType : public Type::TypeBase<TokenType, Type, TypeStorage> {
 public:
  using Base::Base;
  static constexpr ::mlir::StringLiteral name = "tf_executor.token";
};

}  // namespace tf_executor
}  // namespace mlir

// Declares the operations for this dialect using the generated header.

func.func @testEmptyRegions() {
  tf_executor.graph {
    tf_executor.island() {
      "tf_device.parallel_execute"() ({
        tf_device.return
      }, {
        tf_device.return
      }) {} : () -> ()
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK:      [[ISLAND_0_CTRL:%.+]] = tf_executor.island {
// CHECK:        tf_executor.yield {_parallel_execution_ids = "p0:0"}

        "tf_executor.yield"(%0) : (tensor<f32>) -> ()
      }) : (!tf_executor.control) -> (tensor<f32>, !tf_executor.control)
      %outputs_6, %control_7 = "tf_executor.island"() ({
        %0 = "tf.Identity"(%outputs_0#1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
        "tf_executor.yield"(%0) : (tensor<*xi32>) -> ()
      }) : () -> (tensor<*xi32>, !tf_executor.control)
      %outputs_8, %control_9 = "tf_executor.island"(%control_7) ({

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

  func.func @main() {
    tf_executor.graph {
      %0 = tf_executor.island wraps "tf.PartitionedCall"() {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : () -> ()
      tf_executor.fetch
    }
    func.return
  }
  func.func @foo() {
    tf_executor.graph {
      %0:2 = tf_executor.island {

      %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor<?xf32>) -> tensor<?xf32>

      // CHECK: tf_executor.yield %[[CALL_OUTPUT]]
      tf_executor.yield %3 : tensor<?xf32>
    }
    tf_executor.fetch %1#0 : tensor<?xf32>
  }
  func.return %0 : tensor<?xf32>
}

func.func @_func(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  func.return %arg0 : tensor<?xf32>
}

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"

#define DEBUG_TYPE "tf-functional-to-executor"

namespace mlir {

namespace {
// This pass converts mlir functions consisting of mlir ops into a tf_executor
// dialect as a single island.
// Result like so:
//   func @my_fn(%argi...) -> (result_t) {
//     %results:[[n_args]] = tf_executor.graph {