// RUN: tf-opt -tfl-reduce-while %s | FileCheck %s

// The original func we want to optimize is:
//
// func increase_3rd_operand_3_times():
//   S = (1, 0, 0)
//   whlie (S[2] < 3) {
//     s0 = S[0] * 2
//     s1 = S[0] + S[1]
//     s2 = S[2] + 1
//     S = (s0, s1, s2)
//   }
//   return S[2]
// }
//
// Since only S[2] is returned and the computation of final S[2] does not depend
// on S[0] and S[1]. The func can be optimized to

fun test() {
    var x = 0
    <expr>consume(x)
    while (cond()) {
        consume(++x)
    }</expr>
}

fun cond(): Boolean = true

// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main() {
  // CHECK:      name: "while/Merge"
  // CHECK-NEXT: op: "Merge"
  // CHECK-NEXT: input: "while/Enter"
  // CHECK-NEXT: input: "while/NextIteration"
  // CHECK:      name: "while/NextIteration"
  // CHECK-NEXT: op: "NextIteration"
  // CHECK-NEXT: input: "while/Add"
  tf_executor.graph {

    }
  }
}
node {
  name: "while/Less"
  op: "Less"
  input: "while/Merge"
  input: "while/Less/y"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "while/LoopCond"
  op: "LoopCond"
  input: "while/Less"
}
node {
  name: "while/Switch"
  op: "Switch"
  input: "while/Merge"
  input: "while/LoopCond"
  attr {
    key: "T"

    %8:2 = tf_executor.Exit %7#1 : tensor<*x!tf_type.int32ref> {device = "", T = "tfdtype$DT_INT32"} loc("while/Exit")
    %10:2 = tf_executor.island(%7#2) wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_INT32", value = dense<1> : tensor<i32>} : () -> tensor<i32> loc("while/Add/y")

  }
  func.return %0#0, %0#1 : tensor<*xi32>, tensor<*xf32>
}

// Verify that While op is mapped to TensorFlow StatelessWhile op if the
// is_stateless attribute is present and otherwise it is mapped to TensorFlow
// While op. In both cases, the additional attribute should be dropped.

// CHECK: name: "StatefulWhile"
// CHECK-NOT: name:
// CHECK: op: "While"
// CHECK-NOT: is_stateless
// CHECK-NOT: shape_invariant
// CHECK:  attr {

// RUN: tf-opt --tfl-legalize-tf-while %s -o - | FileCheck %s
// RUN: tf-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline='default-pipeline=''' | FileCheck %s --check-prefix=INLINE
// RUN: tf-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline | FileCheck %s --check-prefix=CANON

# Verify that TensorFlow While and StatelessWhile ops are mapped to the
# composite While op in MLIR with is_stateless attribute set accordingly to
# distinguish between them.

# CHECK-DAG: "tf.While"{{.*}} is_stateless = false{{.*}} loc(fused["While:", "StatefulWhile"])
# CHECK-DAG: "tf.While"{{.*}} is_stateless = true{{.*}} loc(fused["StatelessWhile:", "StatelessWhile"])

Dmitrii Gridin <******@****.***> 1716981664 +0200

// only check the message content here.

// CHECK-NOT: expected control flow function {{.*}} to have exactly 1 use, found 2

"builtin.module"() ({
  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):