%3 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<2> : tensor<1xi64>}> : (tensor<256xi32>) -> tensor<4x32x256xi32>
  %4:2 = "mhlo.reduce"(%arg0, %3, %0, %1) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<i32>, %arg3: tensor<f32>, %arg4: tensor<i32>):
    %7 = "mhlo.compare"(%arg1, %arg3) {comparison_direction = #mhlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>

  // CHECK-NOT: f32
  // CHECK: stablehlo.reduce_window
  %0 = stablehlo.constant dense<0.0> : tensor<f32>
  %1 = "stablehlo.reduce_window"(%arg0, %0) ({
    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
      %2 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
      stablehlo.return %2 : tensor<f32>

  %3 = "mhlo.reduce_window"(%arg0, %2) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
    %7 = mhlo.add %arg1, %arg2 : tensor<f32>
    mhlo.return %7 : tensor<f32>
  }) {window_dimensions = dense<[1, 1, 3, 3]> : tensor<4xi64>} : (tensor<1x3x6x6xf32>, tensor<f32>) -> tensor<1x3x4x4xf32>
  %4 = "mhlo.reduce_window"(%1, %2) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
    %7 = mhlo.add %arg1, %arg2 : tensor<f32>

  %while:2 = "tf.WhileRegion"(%ta#1, %size) ({
  // CHECK: ^bb0(%[[BARG0:.*]]: tensor<f32>, %[[BARG1:.*]]: tensor<i32>):
  ^bb0(%barg0: tensor<f32>, %barg1: tensor<i32>):
    // CHECK: %[[PRED:.*]] = "tf._SomeOp"(%[[BARG1]])
    // CHECK: "tf.Yield"(%[[PRED]])
    %pred = "tf._SomeOp"(%barg1) : (tensor<i32>) -> tensor<i1>
    "tf.Yield" (%pred) : (tensor<i1>) -> ()
  }, {
  // CHECK: ^bb0(%[[BARG0:.*]]: tensor<f32>, %[[BARG1:.*]]: tensor<i32>):

  } : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>
  %1 = "tf.WhileRegion"(%0) ({
  ^bb0(%carg0: tensor<?x!tf_type.qint8>):
    %check = "tf.Equal"(%arg1, %one) : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "tf.Yield"(%check) : (tensor<i1>) -> ()
    }, {
  ^bb0(%barg0: tensor<?x!tf_type.qint8>):
    %id = "tf.Identity"(%barg0) : (tensor<?x!tf_type.qint8>) -> tensor<?x!tf_type.qint8>

package ssa

import "cmd/internal/src"

// branchelim tries to eliminate branches by
// generating CondSelect instructions.
//
// Search for basic blocks that look like
//
//	bb0            bb0
//	 | \          /   \
//	 | bb1  or  bb1   bb2    <- trivial if/else blocks
//	 | /          \   /
//	bb2            bb3
//
// where the intermediate blocks are mostly empty (with no side-effects);

  // CHECK: ^bb0(%[[VALUE_A:.*]]: tensor<f32>, %[[VALUE_B:.*]]: tensor<f32>):
  // CHECK: %[[SELECT_RESULT:.*]] = mhlo.compare GE, %[[VALUE_A]], %[[VALUE_B]], NOTYPE : (tensor<f32>, tensor<f32>) -> tensor<i1>
  // CHECK: mhlo.return %[[SELECT_RESULT]] : tensor<i1>
  // CHECK: },  {
  // CHECK: ^bb0(%[[VALUE_A:.*]]: tensor<f32>, %[[VALUE_B:.*]]: tensor<f32>):

func.func @conv(tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>, tensor<256x3x32x32xf32>) -> (tensor<256x8x7x16xf32>, tensor<256x16x32x32xf32>, tensor<256x8x6x16xf32>, tensor<256x32x32x16xf32>, tensor<256x32x32x16xf32>) {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<3x3x3x16xf32>, %arg2: tensor<256x3x32x32xf32>) :
   // OK

      %cst = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
      %1:2 = "tf.WhileRegion"(%cst, %arg1) ({
      ^bb0(%arg2: tensor<i1>, %arg3: tensor<i64>):
        "tf.Yield"(%arg2) : (tensor<i1>) -> ()
      }, {
      ^bb0(%arg2: tensor<i1>, %arg3: tensor<i64>):

  %1:2 = "tf.WhileRegion"(%cst_0, %0) <{is_stateless = false}> ({
      ^bb0(%carg0: tensor<i32>, %carg1: tensor<!tf_type.variant>):
       %cst_2 = arith.constant dense<2> : tensor<i32>
       %1 = "tf.Less"(%carg0, %cst_2) : (tensor<i32>, tensor<i32>) -> tensor<i1>
       "tf.Yield"(%1) : (tensor<i1>) -> ()

      // verify condition types
      // CHECK: ^bb0(%[[CARG0:.*]]: tensor<i32>, %[[CARG1:.*]]: tensor<*xf32>):