// -----

// Tests PartitionedCall/StatefulPartitionedCall.

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<i1>) -> () {
  %max_size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NOT: tf.Stack
  %stack = "tf.StackV2"(%max_size) {elem_type = f32, stack_name = "s"} : (tensor<i32>) -> tensor<!tf_type.resource>
  // CHECK: "tf.StatefulPartitionedCall"
  // CHECK-SAME: f = @callee_stack_decomposed

    %0 = tf_executor.graph {

  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
  // CHECK: "tf.StatefulPartitionedCall"([[handle]])
  %x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)

  // CHECK-NOT: tf.VarHandleOp
  // CHECK:  "tf._TfrtGetResource"()
  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
  // CHECK: tf.StatefulPartitionedCall
  %x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)

  // When add_edges_to_output_of_downstream_nodes is true, the output edges of
  // the xla_launch_node's immediate downstream nodes would be attached to the
  // generated xla node. For example, if the original graph is
  // StatefulPartitionedCall{_xla_compile_id=1} -> XlaClusterOutput -> NodeA
  // The output graph of this function would look like the following when
  // add_edges_to_output_of_downstream_nodes is true:
  // XlaLaunch -> NodeA

    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
    }

// -----

// This test makes sure functions with tf._noinline=true is not inlined.

module {
  func.func @stateful_partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> (tensor<1x2x2x3xf32>) {
    %0 = "tf.StatefulPartitionedCall"(%arg0) <{
      config = "", config_proto = "", executor_type = "", f = @some_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>

  // CHECK: func @f(%arg0
  func.func @f(%arg0: tensor<!tf_type.string> {tf_saved_model.bound_input = @v})
  attributes {tf_saved_model.exported_names = ["f"]} {
    "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee} : (tensor<!tf_type.string>) -> ()
    func.return
  }

  func.func private @f_callee(%arg0: tensor<!tf_type.string>) {
    func.return

def XlaClusterFormationPass : Pass<"tf-xla-cluster-formation", "ModuleOp"> {
  let summary = "Encapsulate partitioned calls within a Cluster op";
  let description = [{
    This pass clusters `tf.PartitionedCall` and `tf.StatefulPartitionedCall`
    with `_xla_compile_device_type` attribute into a `tf_device.cluster`.
    Notice this pass will only rewrite the outermost call if there are nested

  %tl = "tf.EmptyTensorList"(%elem_shape, %max_size) : (tensor<0xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<f32>>>
  // CHECK: "tf.StatefulPartitionedCall"(%[[INIT]],
  // CHECK-SAME: f = @callee_tensorlist_decomposed
  %call = "tf.StatefulPartitionedCall"(%tl, %arg0) {f = @callee, config = "", config_proto = "", executor_type = ""}