// RUN: iree-run-mlir --iree-input-type=mhlo --iree-hal-target-backends=vmvx %s | FileCheck %s
// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo --iree-hal-target-backends=vulkan-spirv %s | FileCheck %s)

// CHECK-LABEL: EXEC @collatz
func.func @collatz() -> tensor<f32> {
  %arg0 = util.unfoldable_constant dense<178.0> : tensor<f32>
  %0 = mhlo.constant dense<1.0> : tensor<f32>
  %1 = mhlo.constant dense<3.0> : tensor<f32>
  %2 = mhlo.constant dense<2.0> : tensor<f32>
  %3 = mhlo.constant dense<0.0> : tensor<f32>
  cf.br ^bb1(%3, %arg0 : tensor<f32>, tensor<f32>)
^bb1(%4: tensor<f32>, %5: tensor<f32>):
  %6 = "mhlo.compare"(%5, %0) {comparison_direction = #mhlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
  %7 = tensor.extract %6[] : tensor<i1>
  cf.cond_br %7, ^bb2(%4, %5 : tensor<f32>, tensor<f32>), ^bb6(%4 : tensor<f32>)
^bb2(%8: tensor<f32>, %9: tensor<f32>):
  %10 = mhlo.add %8, %0 : tensor<f32>
  %11 = mhlo.remainder %9, %2 : tensor<f32>
  %12 = "mhlo.compare"(%11, %3) {comparison_direction = #mhlo<comparison_direction NE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
  %13 = "mhlo.compare"(%2, %3) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
  %14 = "mhlo.compare"(%11, %3) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
  %15 = "mhlo.compare"(%13, %14) {comparison_direction = #mhlo<comparison_direction NE>} : (tensor<i1>, tensor<i1>) -> tensor<i1>
  %16 = mhlo.and %12, %15 : tensor<i1>
  %17 = mhlo.add %11, %2 : tensor<f32>
  %18 = "mhlo.select"(%16, %17, %11) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
  %19 = "mhlo.compare"(%18, %3) {comparison_direction = #mhlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
  %20 = tensor.extract %19[] : tensor<i1>
  cf.cond_br %20, ^bb3, ^bb4
^bb3: // pred: ^bb2
  %21 = mhlo.multiply %9, %1 : tensor<f32>
  %22 = mhlo.add %21, %0 : tensor<f32>
  cf.br ^bb5(%22 : tensor<f32>)
^bb4: // pred: ^bb2
  %23 = mhlo.divide %9, %2 : tensor<f32>
  cf.br ^bb5(%23 : tensor<f32>)
^bb5(%24: tensor<f32>): // 2 preds: ^bb3, ^bb4
  cf.br ^bb1(%10, %24 : tensor<f32>, tensor<f32>)
^bb6(%25: tensor<f32>): // pred: ^bb1
  return %25 : tensor<f32>
}
// CHECK: f32=31