TheRock/compiler/amd-llvm/mlir/test/Dialect/ControlFlow/canonicalize.mlir

// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='builtin.module(func.func(canonicalize{test-convergence}))' -split-input-file | FileCheck --dump-input-context 20 %s

/// Test the folding of BranchOp.

// CHECK-LABEL: func @br_folding(
func.func @br_folding() -> i32 {
  // CHECK-NEXT: %[[CST:.*]] = arith.constant 0 : i32
  // CHECK-NEXT: return %[[CST]] : i32
  %c0_i32 = arith.constant 0 : i32
  cf.br ^bb1(%c0_i32 : i32)
^bb1(%x : i32):
  return %x : i32
}

/// Test that pass-through successors of BranchOp get folded.

// CHECK-LABEL: func @br_passthrough(
// CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
func.func @br_passthrough(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
  "foo.switch"() [^bb1, ^bb2, ^bb3] : () -> ()

^bb1:
  // CHECK: ^bb1:
  // CHECK-NEXT: cf.br ^bb3(%[[ARG0]], %[[ARG1]] : i32, i32)

  cf.br ^bb2(%arg0 : i32)

^bb2(%arg2 : i32):
  cf.br ^bb3(%arg2, %arg1 : i32, i32)

^bb3(%arg4 : i32, %arg5 : i32):
  return %arg4, %arg5 : i32, i32
}

/// Test the folding of CondBranchOp with a constant condition.

// CHECK-LABEL: func @cond_br_folding(
func.func @cond_br_folding(%cond : i1, %a : i32) {
  // CHECK-NEXT: return

  %false_cond = arith.constant false
  %true_cond = arith.constant true
  cf.cond_br %cond, ^bb1, ^bb2(%a : i32)

^bb1:
  cf.cond_br %true_cond, ^bb3, ^bb2(%a : i32)

^bb2(%x : i32):
  cf.cond_br %false_cond, ^bb2(%x : i32), ^bb3

^bb3:
  return
}

/// Test the folding of CondBranchOp when the successors are identical.

// CHECK-LABEL: func @cond_br_same_successor(
func.func @cond_br_same_successor(%cond : i1, %a : i32) {
  // CHECK-NEXT: return

  cf.cond_br %cond, ^bb1(%a : i32), ^bb1(%a : i32)

^bb1(%result : i32):
  return
}

/// Test the folding of CondBranchOp when the successors are identical, but the
/// arguments are different.

// CHECK-LABEL: func @cond_br_same_successor_insert_select(
// CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
// CHECK-SAME: %[[ARG2:.*]]: tensor<2xi32>, %[[ARG3:.*]]: tensor<2xi32>
func.func @cond_br_same_successor_insert_select(
      %cond : i1, %a : i32, %b : i32, %c : tensor<2xi32>, %d : tensor<2xi32>
    ) -> (i32, tensor<2xi32>)  {
  // CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG1]]
  // CHECK: %[[RES2:.*]] = arith.select %[[COND]], %[[ARG2]], %[[ARG3]]
  // CHECK: return %[[RES]], %[[RES2]]

  cf.cond_br %cond, ^bb1(%a, %c : i32, tensor<2xi32>), ^bb1(%b, %d : i32, tensor<2xi32>)

^bb1(%result : i32, %result2 : tensor<2xi32>):
  return %result, %result2 : i32, tensor<2xi32>
}

/// Test the compound folding of BranchOp and CondBranchOp.

// CHECK-LABEL: func @cond_br_and_br_folding(
func.func @cond_br_and_br_folding(%a : i32) {
  // CHECK-NEXT: return

  %false_cond = arith.constant false
  %true_cond = arith.constant true
  cf.cond_br %true_cond, ^bb2, ^bb1(%a : i32)

^bb1(%x : i32):
  cf.cond_br %false_cond, ^bb1(%x : i32), ^bb2

^bb2:
  return
}

/// Test that pass-through successors of CondBranchOp get folded.

// Test that the weights are preserved:
// CHECK-LABEL:   func.func @cond_br_passthrough_weights(
// CHECK-SAME:      %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i1) -> i32 {
func.func @cond_br_passthrough_weights(%arg0 : i32, %arg1 : i32, %cond : i1) -> i32 {
// CHECK:           cf.cond_br %[[ARG2]] weights([30, 70]), ^bb1, ^bb2
// CHECK:         ^bb1:
// CHECK:           return %[[ARG0]] : i32
// CHECK:         ^bb2:
// CHECK:           return %[[ARG1]] : i32
// CHECK:         }
  cf.cond_br %cond weights([30,70]), ^bb1, ^bb3

^bb1:
  cf.br ^bb2

^bb3:
  cf.br ^bb4

^bb2:
  return %arg0 : i32

^bb4:
  return %arg1 : i32
}

// CHECK-LABEL: func @cond_br_passthrough(
// CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[COND:.*]]: i1
func.func @cond_br_passthrough(%arg0 : i32, %arg1 : i32, %arg2 : i32, %cond : i1) -> (i32, i32) {
  // CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG2]]
  // CHECK: %[[RES2:.*]] = arith.select %[[COND]], %[[ARG1]], %[[ARG2]]
  // CHECK: return %[[RES]], %[[RES2]]

  cf.cond_br %cond, ^bb1(%arg0 : i32), ^bb2(%arg2, %arg2 : i32, i32)

^bb1(%arg3: i32):
  cf.br ^bb2(%arg3, %arg1 : i32, i32)

^bb2(%arg4: i32, %arg5: i32):
  return %arg4, %arg5 : i32, i32
}

/// Test the failure modes of collapsing CondBranchOp pass-throughs successors.

// CHECK-LABEL: func @cond_br_pass_through_fail(
func.func @cond_br_pass_through_fail(%cond : i1) {
  // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2

  cf.cond_br %cond, ^bb1, ^bb2

^bb1:
  // CHECK: ^bb1:
  // CHECK: "foo.op"
  // CHECK: cf.br ^bb2

  // Successors can't be collapsed if they contain other operations.
  "foo.op"() : () -> ()
  cf.br ^bb2

^bb2:
  return
}


/// Test the folding of SwitchOp

// CHECK-LABEL: func @switch_only_default(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
func.func @switch_only_default(%flag : i32, %caseOperand0 : f32) {
  // add predecessors for all blocks to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2] : () -> ()
  ^bb1:
    // CHECK-NOT: cf.switch
    // CHECK: cf.br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
    cf.switch %flag : i32, [
      default: ^bb2(%caseOperand0 : f32)
    ]
  // CHECK: ^[[BB2]]({{.*}}):
  ^bb2(%bb2Arg : f32):
    // CHECK-NEXT: "foo.bb2Terminator"
    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
}


// CHECK-LABEL: func @switch_case_matching_default(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
func.func @switch_case_matching_default(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
  // add predecessors for all blocks to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb3] : () -> ()
  ^bb1:
    // CHECK: cf.switch %[[FLAG]]
    // CHECK-NEXT:   default: ^[[BB1:.+]](%[[CASE_OPERAND_0]] : f32)
    // CHECK-NEXT:   10: ^[[BB2:.+]](%[[CASE_OPERAND_1]] : f32)
    // CHECK-NEXT: ]
    cf.switch %flag : i32, [
      default: ^bb2(%caseOperand0 : f32),
      42: ^bb2(%caseOperand0 : f32),
      10: ^bb3(%caseOperand1 : f32),
      17: ^bb2(%caseOperand0 : f32)
    ]
  ^bb2(%bb2Arg : f32):
    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
  ^bb3(%bb3Arg : f32):
    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
}


// CHECK-LABEL: func @switch_on_const_no_match(
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
func.func @switch_on_const_no_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
  // add predecessors for all blocks to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
  ^bb1:
    // CHECK-NOT: cf.switch
    // CHECK: cf.br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
    %c0_i32 = arith.constant 0 : i32
    cf.switch %c0_i32 : i32, [
      default: ^bb2(%caseOperand0 : f32),
      -1: ^bb3(%caseOperand1 : f32),
      1: ^bb4(%caseOperand2 : f32)
    ]
  // CHECK: ^[[BB2]]({{.*}}):
  // CHECK-NEXT: "foo.bb2Terminator"
  ^bb2(%bb2Arg : f32):
    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
  ^bb3(%bb3Arg : f32):
    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
}

// CHECK-LABEL: func @switch_on_const_with_match(
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
func.func @switch_on_const_with_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
  // add predecessors for all blocks to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
  ^bb1:
    // CHECK-NOT: cf.switch
    // CHECK: cf.br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
    %c0_i32 = arith.constant 1 : i32
    cf.switch %c0_i32 : i32, [
      default: ^bb2(%caseOperand0 : f32),
      -1: ^bb3(%caseOperand1 : f32),
      1: ^bb4(%caseOperand2 : f32)
    ]
  ^bb2(%bb2Arg : f32):
    "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
  ^bb3(%bb3Arg : f32):
    "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
  // CHECK: ^[[BB4]]({{.*}}):
  // CHECK-NEXT: "foo.bb4Terminator"
  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
}

// CHECK-LABEL: func @switch_passthrough(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_3:[a-zA-Z0-9_]+]]
func.func @switch_passthrough(%flag : i32,
                         %caseOperand0 : f32,
                         %caseOperand1 : f32,
                         %caseOperand2 : f32,
                         %caseOperand3 : f32) {
  // add predecessors for all blocks to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4, ^bb5, ^bb6] : () -> ()

  ^bb1:
  //      CHECK: cf.switch %[[FLAG]]
  // CHECK-NEXT:   default: ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
  // CHECK-NEXT:   43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
  // CHECK-NEXT:   44: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
  // CHECK-NEXT: ]
    cf.switch %flag : i32, [
      default: ^bb2(%caseOperand0 : f32),
      43: ^bb3(%caseOperand1 : f32),
      44: ^bb4(%caseOperand2 : f32)
    ]
  ^bb2(%bb2Arg : f32):
    cf.br ^bb5(%bb2Arg : f32)
  ^bb3(%bb3Arg : f32):
    cf.br ^bb6(%bb3Arg : f32)
  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()

  // CHECK: ^[[BB5]]({{.*}}):
  // CHECK-NEXT: "foo.bb5Terminator"
  ^bb5(%bb5Arg : f32):
    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()

  // CHECK: ^[[BB6]]({{.*}}):
  // CHECK-NEXT: "foo.bb6Terminator"
  ^bb6(%bb6Arg : f32):
    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
}

// CHECK-LABEL: func @switch_from_switch_with_same_value_with_match(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
func.func @switch_from_switch_with_same_value_with_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5] : () -> ()

  ^bb1:
    // CHECK: cf.switch %[[FLAG]]
    cf.switch %flag : i32, [
      default: ^bb2,
      42: ^bb3
    ]

  ^bb2:
    "foo.bb2Terminator"() : () -> ()
  ^bb3:
    // prevent this block from being simplified away
    "foo.op"() : () -> ()
    // CHECK-NOT: cf.switch %[[FLAG]]
    // CHECK: cf.br ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
    cf.switch %flag : i32, [
      default: ^bb4(%caseOperand0 : f32),
      42: ^bb5(%caseOperand1 : f32)
    ]

  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()

  // CHECK: ^[[BB5]]({{.*}}):
  // CHECK-NEXT: "foo.bb5Terminator"
  ^bb5(%bb5Arg : f32):
    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
}

// CHECK-LABEL: func @switch_from_switch_with_same_value_no_match(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
func.func @switch_from_switch_with_same_value_no_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()

  ^bb1:
    // CHECK: cf.switch %[[FLAG]]
    cf.switch %flag : i32, [
      default: ^bb2,
      42: ^bb3
    ]

  ^bb2:
    "foo.bb2Terminator"() : () -> ()
  ^bb3:
    "foo.op"() : () -> ()
    // CHECK-NOT: cf.switch %[[FLAG]]
    // CHECK: cf.br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
    cf.switch %flag : i32, [
      default: ^bb4(%caseOperand0 : f32),
      0: ^bb5(%caseOperand1 : f32),
      43: ^bb6(%caseOperand2 : f32)
    ]

  // CHECK: ^[[BB4]]({{.*}})
  // CHECK-NEXT: "foo.bb4Terminator"
  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()

  ^bb5(%bb5Arg : f32):
    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()

  ^bb6(%bb6Arg : f32):
    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
}

// CHECK-LABEL: func @switch_from_switch_default_with_same_value(
// CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
func.func @switch_from_switch_default_with_same_value(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
  // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
  "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()

  ^bb1:
    // CHECK: cf.switch %[[FLAG]]
    cf.switch %flag : i32, [
      default: ^bb3,
      42: ^bb2
    ]

  ^bb2:
    "foo.bb2Terminator"() : () -> ()
  ^bb3:
    "foo.op"() : () -> ()
    // CHECK: cf.switch %[[FLAG]]
    // CHECK-NEXT: default: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
    // CHECK-NEXT: 43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
    // CHECK-NOT: 42
    cf.switch %flag : i32, [
      default: ^bb4(%caseOperand0 : f32),
      42: ^bb5(%caseOperand1 : f32),
      43: ^bb6(%caseOperand2 : f32)
    ]

  // CHECK: ^[[BB4]]({{.*}}):
  // CHECK-NEXT: "foo.bb4Terminator"
  ^bb4(%bb4Arg : f32):
    "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()

  ^bb5(%bb5Arg : f32):
    "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()

  // CHECK: ^[[BB6]]({{.*}}):
  // CHECK-NEXT: "foo.bb6Terminator"
  ^bb6(%bb6Arg : f32):
    "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
}

/// Test folding conditional branches that are successors of conditional
/// branches with the same condition.

// CHECK-LABEL: func @cond_br_from_cond_br_with_same_condition
func.func @cond_br_from_cond_br_with_same_condition(%cond : i1) {
  // CHECK:   cf.cond_br %{{.*}}, ^bb1, ^bb2
  // CHECK: ^bb1:
  // CHECK:   return

  cf.cond_br %cond, ^bb1, ^bb2

^bb1:
  cf.cond_br %cond, ^bb3, ^bb2

^bb2:
  "foo.terminator"() : () -> ()

^bb3:
  return
}

// -----

// Erase assertion if condition is known to be true at compile time.
// CHECK-LABEL: @assert_true
func.func @assert_true() {
  // CHECK-NOT: cf.assert
  %true = arith.constant true
  cf.assert %true, "Computer says no"
  return
}

// -----

// Keep assertion if condition unknown at compile time.
// CHECK-LABEL: @cf.assert
// CHECK-SAME:  (%[[ARG:.*]]: i1)
func.func @cf.assert(%arg : i1) {
  // CHECK: cf.assert %[[ARG]], "Computer says no"
  cf.assert %arg, "Computer says no"
  return
}

// -----

// CHECK-LABEL: @branchCondProp
//       CHECK:       %[[trueval:.+]] = arith.constant true
//       CHECK:       %[[falseval:.+]] = arith.constant false
//       CHECK:       "test.consumer1"(%[[trueval]]) : (i1) -> ()
//       CHECK:       "test.consumer2"(%[[falseval]]) : (i1) -> ()
func.func @branchCondProp(%arg0: i1) {
  cf.cond_br %arg0, ^trueB, ^falseB

^trueB:
  "test.consumer1"(%arg0) : (i1) -> ()
  cf.br ^exit

^falseB:
  "test.consumer2"(%arg0) : (i1) -> ()
  cf.br ^exit

^exit:
  return
}

// -----

/// Test that control-flow cycles are not simplified infinitely.

// CHECK-LABEL:   @cycle_2_blocks
// CHECK:           cf.br ^bb1
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb1
func.func @cycle_2_blocks() {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb1
}

// CHECK-LABEL:   @no_cycle_2_blocks
// CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
// CHECK:           return %[[VAL_0]] : i32
func.func @no_cycle_2_blocks() -> i32 {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb3
^bb3:
  %ret = arith.constant 1 : i32
  return %ret : i32
}

// CHECK-LABEL:   @cycle_4_blocks
// CHECK:           cf.br ^bb1
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb1
func.func @cycle_4_blocks() {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb3
^bb3:
  cf.br ^bb4
^bb4:
  cf.br ^bb1
}

// CHECK-LABEL:   @no_cycle_4_blocks
// CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
// CHECK:           return %[[VAL_0]] : i32
func.func @no_cycle_4_blocks() -> i32 {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb3
^bb3:
  cf.br ^bb4
^bb4:
  cf.br ^bb5
^bb5:
  %ret = arith.constant 1 : i32
  return %ret : i32
}

// CHECK-LABEL:   @delayed_3_cycle
// CHECK:           cf.br ^bb1
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb1
func.func @delayed_3_cycle() {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb3
^bb3:
  cf.br ^bb4
^bb4:
  cf.br ^bb5
^bb5:
  cf.br ^bb3
}

// CHECK-LABEL:   @cycle_1_block
// CHECK:           cf.br ^bb1
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb1
func.func @cycle_1_block() {
  cf.br ^bb1
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb2
}

// CHECK-LABEL:   @unsimplified_cycle_1
// CHECK-SAME:      %[[ARG0:.*]]: i1) {
// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb2
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb2
// CHECK:         ^bb2:
// CHECK:           cf.br ^bb3
// CHECK:         ^bb3:
// CHECK:           cf.br ^bb3
func.func @unsimplified_cycle_1(%c : i1) {
  cf.cond_br %c, ^bb1, ^bb2
^bb1:
  cf.br ^bb2
^bb2:
  cf.br ^bb3
^bb3:
  cf.br ^bb4
^bb4:
  cf.br ^bb3
}

// Make sure we terminate when other cf passes can't help us.

// CHECK-LABEL:   @unsimplified_cycle_2
// CHECK-SAME:      %[[ARG0:.*]]: i1) {
// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb3
// CHECK:         ^bb1:
// CHECK:           cf.br ^bb2 {A}
// CHECK:         ^bb2:
// CHECK:           cf.br ^bb2 {E}
// CHECK:         ^bb3:
// CHECK:           cf.br ^bb1
func.func @unsimplified_cycle_2(%c : i1) {
  cf.cond_br %c, ^bb6, ^bb7
^bb6:
  cf.br ^bb5 {F}
^bb5:
  cf.br ^bb1 {A}
^bb1:
  cf.br ^bb2 {B}
^bb2:
  cf.br ^bb3 {C}
^bb3:
  cf.br ^bb4 {D}
^bb4:
  cf.br ^bb1 {E}
^bb7:
  cf.br ^bb6
}

// CHECK-LABEL: @drop_unreachable_branch_1
//  CHECK-NEXT:   "test.foo"() : () -> ()
//  CHECK-NEXT:   return
func.func @drop_unreachable_branch_1(%c: i1) {
  cf.cond_br %c, ^bb1, ^bb2
^bb1:
  "test.foo"() : () -> ()
  return
^bb2:
  ub.unreachable
}

// CHECK-LABEL: @drop_unreachable_branch_2
//  CHECK-NEXT:   ub.unreachable
func.func @drop_unreachable_branch_2(%c: i1) {
  cf.cond_br %c, ^bb1, ^bb2
^bb1:
  ub.unreachable
^bb2:
  ub.unreachable
}

// -----

// Regression test for https://github.com/llvm/llvm-project/issues/126213:
// simplifyBrToBlockWithSinglePred must not merge a block into its sole
// predecessor when a branch operand is itself a block argument of the
// destination — replaceAllUsesWith would be a no-op for that argument,
// leaving the block argument live after the block is erased and crashing.
//
// ^guard is merged into the entry block (single predecessor), replacing %cond
// with %true and folding the cond_br. This kills the external entry to
// ^header, leaving ^body as its only predecessor. Without the fix,
// simplifyBrToBlockWithSinglePred would attempt to merge ^header into ^body
// despite %keep being a block argument of ^header, triggering the crash.

// CHECK-LABEL: @no_merge_self_arg_loop
//  CHECK-NEXT:   %[[TRUE:.*]] = arith.constant true
//  CHECK-NEXT:   return %[[TRUE]]
func.func @no_merge_self_arg_loop(%step: i1) -> i1 {
  %true = arith.constant true
  cf.br ^guard(%true : i1)
^guard(%cond: i1):
  cf.cond_br %cond, ^exit(%true : i1), ^header(%step, %true : i1, i1)
^header(%iter: i1, %keep: i1):
  cf.cond_br %iter, ^body, ^exit(%keep : i1)
^body:
  cf.br ^header(%step, %keep : i1, i1)
^exit(%result: i1):
  return %result : i1
}

// Verify that block arguments are replaced with a uniform incoming value
// when all predecessors pass the same SSA value

// CHECK-LABEL: func @fold_uniform_branch_block_arg
// CHECK-SAME: %[[COND:.*]]: i1, %[[C:.*]]: i32
func.func @fold_uniform_branch_block_arg(%cond: i1, %c: i32) -> i32 {
  cf.cond_br %cond, ^bb1, ^bb2
^bb1:
  "foo.op"() : () -> ()
  cf.br ^bb3(%c : i32)
^bb2:
  "foo.op"() : () -> ()
  cf.br ^bb3(%c : i32)
^bb3(%arg0: i32):
  // CHECK: ^bb3:
  // CHECK: return %[[C]]
  return %arg0 : i32
}

// Verify that block arguments are not folded when incoming values differ
// across predecessors.

// CHECK-LABEL: func @no_fold_non_uniform_block_arg
// CHECK-SAME: %[[COND:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
func.func @no_fold_non_uniform_block_arg(%cond: i1, %a: i32, %b: i32) -> i32 {
  cf.cond_br %cond, ^bb1, ^bb2
^bb1:
  "foo.op"() : () -> ()
  cf.br ^bb3(%a : i32)
^bb2:
  "foo.op"() : () -> ()
  cf.br ^bb3(%b : i32)
^bb3(%arg0: i32):
  // CHECK: ^bb3(%[[ARG0:.*]]: i32):
  // CHECK-NEXT: return %[[ARG0]]
  return %arg0 : i32
}

// Verify no folding when the same block appears multiple times as a
// successor with different operands.

// CHECK-LABEL: func @no_fold_same_dest_different_args
func.func @no_fold_same_dest_different_args(%a: i32, %b: i32) -> i32 {
  "test.producing_br"(%a, %b)[^bb1, ^bb1]
    {operandSegmentSizes = array<i32: 1, 1>} : (i32, i32) -> i32
^bb1(%arg0: i32):
  // CHECK: ^bb1(%[[ARG0:.*]]: i32):
  // CHECK-NEXT: return %[[ARG0]]
  return %arg0 : i32
}

// Verify folding when the same block appears multiple times as a
// successor with the same operand.

// CHECK-LABEL: func @fold_same_dest_same_args
// CHECK-SAME: %[[A:.*]]: i32
func.func @fold_same_dest_same_args(%a: i32) -> i32 {
  "test.producing_br"(%a, %a)[^bb1, ^bb1]
      {operandSegmentSizes = array<i32: 1, 1>} : (i32, i32) -> i32
^bb1(%arg0: i32):
  // CHECK: ^bb1:
  // CHECK-NEXT: return %[[A]]
  return %arg0 : i32
}

// Verify no folding when a predecessor has an unknown terminator.

// CHECK-LABEL: func @no_fold_unknown_terminator
func.func @no_fold_unknown_terminator(%a: i32) -> i32 {
  cf.br ^bb1
^bb1:
  "foo.two_successors"()[^bb2, ^bb3] : () -> ()
^bb2:
  // CHECK: ^bb2(%[[ARG0:.*]]: i32):
  cf.br ^bb3(%a : i32)
^bb3(%arg0: i32):
  // CHECK-NEXT: return %[[ARG0]]
  return %arg0 : i32
}

// Verify that unused block arguments are skipped and only used arguments
// with uniform incoming values are folded.

// CHECK-LABEL: func @skip_unused_block_arg
func.func @skip_unused_block_arg(%flag: i32, %a: i32, %b: i32) -> i32 {
  "foo.pred"()[^bb1, ^bb2] : () -> ()
^bb1:
  cf.br ^bb3(%b, %a : i32, i32)
^bb2:
  cf.br ^bb3(%a, %a : i32, i32)
^bb3(%arg0: i32, %arg1: i32):
  "foo.use"(%arg0) : (i32) -> ()
  // CHECK: ^bb3(%[[ARG0:.*]]: i32):
  // CHECK-NEXT: "foo.use"(%[[ARG0]])
  // CHECK-NEXT: return %[[A:.*]]
  return %arg1 : i32
}