// This pass forms `tf_executor.island` per region of
// `tf_device.parallel_execute`.
//
// For example, the following:
//
//  %0 = tf_executor.island {
//    tf_executor.yield
//  }
//  %1:2 = tf_executor.island {
//    %2 = "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
//      tf_executor.yield %2 : tensor<i1>
//  }
//  %3:2 = tf_executor.island(%0) {
//    %4 = "tf_device.parallel_execute"() ({

// CHECK: }

func.func @main() {
  tf_executor.graph {
    %dim:2 = tf_executor.island wraps "tf.Const"() {dtype = "tftype$DT_INT32", value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %input:2 = tf_executor.island wraps "tf.Const"() {dtype = "tftype$DT_INT32", value = dense<1.0> : tensor<4x6xf32>} : () -> tensor<4x6xf32>

      return mlir::WalkResult::interrupt();
    }

    for (mlir::Operation& op : graph.GetBody()) {
      auto island = llvm::dyn_cast<mlir::tf_executor::IslandOp>(op);
      if (!island) continue;

      if (!island.WrapsSingleOp()) {
        island.emitError(kInvalidExecutorGraphMsg)
            << "tf_executor.island must perfectly wrap a single op";
        return mlir::WalkResult::interrupt();
      }
    }

  func.func @device_ordinal_attr_added_multiple_devices() {
    tf_executor.graph {
      %0 = tf_executor.island {
        "tf_device.launch"() ({
          %1 = "tf.opA"() : () -> tensor<!tf_type.string>

// CHECK: %[[OUT_0:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.VarHandleOp"()
// CHECK: %[[OUT_1:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.RestoreV2"
// CHECK: %[[CTL_2:.*]] = tf_executor.island(%[[CTL_1]]) wraps "tf.AssignVariableOp"(%[[OUT_0]], %[[OUT_1]])

// Check that ReadVariableOp now use the same variable op.
// CHECK: %[[OUT_3:.*]], %[[CTL_3:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[OUT_0]])

// -----

    %0 = tf_executor.graph {
      %outputs_0, %control_0 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<501000x128xf32>>>
      %outputs_1, %control_1 = tf_executor.island wraps "tf.Cast"(%outputs_0) {Truncate = false} : (tensor<!tf_type.resource<tensor<501000x128xf32>>>) -> tensor<*x!tf_type.resource>

// CHECK: [[VARIABLE_REG:%.*]], [[VARIABLE_REG_control:%.*]] = tf_executor.island wraps "tf.VarHandleOp"()

// Test for the presence of Identity op between input and function call.
// CHECK: [[IDENTITY_REG:%.*]], [[IDENTITY_REG_control:%.*]] = tf_executor.island wraps "tf.Identity"([[VARIABLE_REG]])

// CHECK: [[CALL_RESULT_REG:%.*]], [[CALL_RESULT_REG_control:%.*]] = tf_executor.island wraps "tf.StatefulPartitionedCall"([[IDENTITY_REG]])

// expected-error@-1 {{'tf_executor.island' op expected 1 or more results}}
    }) : () -> ()
  }
  func.return
}

// -----

// Check that an island region can't be empty.
func.func @invalid_island(%arg0: tensor<*xf32>, %ctl: !tf_executor.control) {
  tf_executor.graph {
    "tf_executor.island"() ({

func.func @main() -> tensor<2x3xi32> {
  %graph = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<5> : tensor<2x3xi32>} : () -> tensor<2x3xi32> loc("Const0")
    %1:2 = tf_executor.island wraps "tf.Const"() {value = dense<4.2> : tensor<4x5xf32>} : () -> tensor<4x5xf32> loc("Const1")

// CHECK-NEXT:       }
// CHECK:          }
// CHECK-NEXT:   }
func.func @bar() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {dtype = "tfdtype$DT_INT32", name = "tf.Const", value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %1:2 = tf_executor.island wraps "tf.Empty"(%0#0) {dtype = "tfdtype$DT_FLOAT", name = "tf.Empty"} : (tensor<i32>) -> tensor<*xf32>
    tf_executor.fetch
  }
  func.return
}