mlir/test/Dialect/ControlFlow/canonicalize.mlir

   1 // 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
   2
   3 /// Test the folding of BranchOp.
   4
   5 // CHECK-LABEL: func @br_folding(
   6 func.func @br_folding() -> i32 {
   7   // CHECK-NEXT: %[[CST:.*]] = arith.constant 0 : i32
   8   // CHECK-NEXT: return %[[CST]] : i32
   9   %c0_i32 = arith.constant 0 : i32
  10   cf.br ^bb1(%c0_i32 : i32)
  11 ^bb1(%x : i32):
  12   return %x : i32
  13 }
  14
  15 /// Test that pass-through successors of BranchOp get folded.
  16
  17 // CHECK-LABEL: func @br_passthrough(
  18 // CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
  19 func.func @br_passthrough(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
  20   "foo.switch"() [^bb1, ^bb2, ^bb3] : () -> ()
  21
  22 ^bb1:
  23   // CHECK: ^bb1:
  24   // CHECK-NEXT: cf.br ^bb3(%[[ARG0]], %[[ARG1]] : i32, i32)
  25
  26   cf.br ^bb2(%arg0 : i32)
  27
  28 ^bb2(%arg2 : i32):
  29   cf.br ^bb3(%arg2, %arg1 : i32, i32)
  30
  31 ^bb3(%arg4 : i32, %arg5 : i32):
  32   return %arg4, %arg5 : i32, i32
  33 }
  34
  35 /// Test the folding of CondBranchOp with a constant condition.
  36
  37 // CHECK-LABEL: func @cond_br_folding(
  38 func.func @cond_br_folding(%cond : i1, %a : i32) {
  39   // CHECK-NEXT: return
  40
  41   %false_cond = arith.constant false
  42   %true_cond = arith.constant true
  43   cf.cond_br %cond, ^bb1, ^bb2(%a : i32)
  44
  45 ^bb1:
  46   cf.cond_br %true_cond, ^bb3, ^bb2(%a : i32)
  47
  48 ^bb2(%x : i32):
  49   cf.cond_br %false_cond, ^bb2(%x : i32), ^bb3
  50
  51 ^bb3:
  52   return
  53 }
  54
  55 /// Test the folding of CondBranchOp when the successors are identical.
  56
  57 // CHECK-LABEL: func @cond_br_same_successor(
  58 func.func @cond_br_same_successor(%cond : i1, %a : i32) {
  59   // CHECK-NEXT: return
  60
  61   cf.cond_br %cond, ^bb1(%a : i32), ^bb1(%a : i32)
  62
  63 ^bb1(%result : i32):
  64   return
  65 }
  66
  67 /// Test the folding of CondBranchOp when the successors are identical, but the
  68 /// arguments are different.
  69
  70 // CHECK-LABEL: func @cond_br_same_successor_insert_select(
  71 // CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
  72 // CHECK-SAME: %[[ARG2:.*]]: tensor<2xi32>, %[[ARG3:.*]]: tensor<2xi32>
  73 func.func @cond_br_same_successor_insert_select(
  74       %cond : i1, %a : i32, %b : i32, %c : tensor<2xi32>, %d : tensor<2xi32>
  75     ) -> (i32, tensor<2xi32>)  {
  76   // CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG1]]
  77   // CHECK: %[[RES2:.*]] = arith.select %[[COND]], %[[ARG2]], %[[ARG3]]
  78   // CHECK: return %[[RES]], %[[RES2]]
  79
  80   cf.cond_br %cond, ^bb1(%a, %c : i32, tensor<2xi32>), ^bb1(%b, %d : i32, tensor<2xi32>)
  81
  82 ^bb1(%result : i32, %result2 : tensor<2xi32>):
  83   return %result, %result2 : i32, tensor<2xi32>
  84 }
  85
  86 /// Test the compound folding of BranchOp and CondBranchOp.
  87
  88 // CHECK-LABEL: func @cond_br_and_br_folding(
  89 func.func @cond_br_and_br_folding(%a : i32) {
  90   // CHECK-NEXT: return
  91
  92   %false_cond = arith.constant false
  93   %true_cond = arith.constant true
  94   cf.cond_br %true_cond, ^bb2, ^bb1(%a : i32)
  95
  96 ^bb1(%x : i32):
  97   cf.cond_br %false_cond, ^bb1(%x : i32), ^bb2
  98
  99 ^bb2:
 100   return
 101 }
 102
 103 /// Test that pass-through successors of CondBranchOp get folded.
 104
 105 // CHECK-LABEL: func @cond_br_passthrough(
 106 // CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[COND:.*]]: i1
 107 func.func @cond_br_passthrough(%arg0 : i32, %arg1 : i32, %arg2 : i32, %cond : i1) -> (i32, i32) {
 108   // CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG2]]
 109   // CHECK: %[[RES2:.*]] = arith.select %[[COND]], %[[ARG1]], %[[ARG2]]
 110   // CHECK: return %[[RES]], %[[RES2]]
 111
 112   cf.cond_br %cond, ^bb1(%arg0 : i32), ^bb2(%arg2, %arg2 : i32, i32)
 113
 114 ^bb1(%arg3: i32):
 115   cf.br ^bb2(%arg3, %arg1 : i32, i32)
 116
 117 ^bb2(%arg4: i32, %arg5: i32):
 118   return %arg4, %arg5 : i32, i32
 119 }
 120
 121 /// Test the failure modes of collapsing CondBranchOp pass-throughs successors.
 122
 123 // CHECK-LABEL: func @cond_br_pass_through_fail(
 124 func.func @cond_br_pass_through_fail(%cond : i1) {
 125   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 126
 127   cf.cond_br %cond, ^bb1, ^bb2
 128
 129 ^bb1:
 130   // CHECK: ^bb1:
 131   // CHECK: "foo.op"
 132   // CHECK: cf.br ^bb2
 133
 134   // Successors can't be collapsed if they contain other operations.
 135   "foo.op"() : () -> ()
 136   cf.br ^bb2
 137
 138 ^bb2:
 139   return
 140 }
 141
 142
 143 /// Test the folding of SwitchOp
 144
 145 // CHECK-LABEL: func @switch_only_default(
 146 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 147 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 148 func.func @switch_only_default(%flag : i32, %caseOperand0 : f32) {
 149   // add predecessors for all blocks to avoid other canonicalizations.
 150   "foo.pred"() [^bb1, ^bb2] : () -> ()
 151   ^bb1:
 152     // CHECK-NOT: cf.switch
 153     // CHECK: cf.br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
 154     cf.switch %flag : i32, [
 155       default: ^bb2(%caseOperand0 : f32)
 156     ]
 157   // CHECK: ^[[BB2]]({{.*}}):
 158   ^bb2(%bb2Arg : f32):
 159     // CHECK-NEXT: "foo.bb2Terminator"
 160     "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
 161 }
 162
 163
 164 // CHECK-LABEL: func @switch_case_matching_default(
 165 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 166 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 167 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 168 func.func @switch_case_matching_default(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
 169   // add predecessors for all blocks to avoid other canonicalizations.
 170   "foo.pred"() [^bb1, ^bb2, ^bb3] : () -> ()
 171   ^bb1:
 172     // CHECK: cf.switch %[[FLAG]]
 173     // CHECK-NEXT:   default: ^[[BB1:.+]](%[[CASE_OPERAND_0]] : f32)
 174     // CHECK-NEXT:   10: ^[[BB2:.+]](%[[CASE_OPERAND_1]] : f32)
 175     // CHECK-NEXT: ]
 176     cf.switch %flag : i32, [
 177       default: ^bb2(%caseOperand0 : f32),
 178       42: ^bb2(%caseOperand0 : f32),
 179       10: ^bb3(%caseOperand1 : f32),
 180       17: ^bb2(%caseOperand0 : f32)
 181     ]
 182   ^bb2(%bb2Arg : f32):
 183     "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
 184   ^bb3(%bb3Arg : f32):
 185     "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
 186 }
 187
 188
 189 // CHECK-LABEL: func @switch_on_const_no_match(
 190 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 191 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 192 // CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
 193 func.func @switch_on_const_no_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
 194   // add predecessors for all blocks to avoid other canonicalizations.
 195   "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
 196   ^bb1:
 197     // CHECK-NOT: cf.switch
 198     // CHECK: cf.br ^[[BB2:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
 199     %c0_i32 = arith.constant 0 : i32
 200     cf.switch %c0_i32 : i32, [
 201       default: ^bb2(%caseOperand0 : f32),
 202       -1: ^bb3(%caseOperand1 : f32),
 203       1: ^bb4(%caseOperand2 : f32)
 204     ]
 205   // CHECK: ^[[BB2]]({{.*}}):
 206   // CHECK-NEXT: "foo.bb2Terminator"
 207   ^bb2(%bb2Arg : f32):
 208     "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
 209   ^bb3(%bb3Arg : f32):
 210     "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
 211   ^bb4(%bb4Arg : f32):
 212     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 213 }
 214
 215 // CHECK-LABEL: func @switch_on_const_with_match(
 216 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 217 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 218 // CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
 219 func.func @switch_on_const_with_match(%caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
 220   // add predecessors for all blocks to avoid other canonicalizations.
 221   "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4] : () -> ()
 222   ^bb1:
 223     // CHECK-NOT: cf.switch
 224     // CHECK: cf.br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
 225     %c0_i32 = arith.constant 1 : i32
 226     cf.switch %c0_i32 : i32, [
 227       default: ^bb2(%caseOperand0 : f32),
 228       -1: ^bb3(%caseOperand1 : f32),
 229       1: ^bb4(%caseOperand2 : f32)
 230     ]
 231   ^bb2(%bb2Arg : f32):
 232     "foo.bb2Terminator"(%bb2Arg) : (f32) -> ()
 233   ^bb3(%bb3Arg : f32):
 234     "foo.bb3Terminator"(%bb3Arg) : (f32) -> ()
 235   // CHECK: ^[[BB4]]({{.*}}):
 236   // CHECK-NEXT: "foo.bb4Terminator"
 237   ^bb4(%bb4Arg : f32):
 238     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 239 }
 240
 241 // CHECK-LABEL: func @switch_passthrough(
 242 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 243 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 244 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 245 // CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
 246 // CHECK-SAME: %[[CASE_OPERAND_3:[a-zA-Z0-9_]+]]
 247 func.func @switch_passthrough(%flag : i32,
 248                          %caseOperand0 : f32,
 249                          %caseOperand1 : f32,
 250                          %caseOperand2 : f32,
 251                          %caseOperand3 : f32) {
 252   // add predecessors for all blocks to avoid other canonicalizations.
 253   "foo.pred"() [^bb1, ^bb2, ^bb3, ^bb4, ^bb5, ^bb6] : () -> ()
 254
 255   ^bb1:
 256   //      CHECK: cf.switch %[[FLAG]]
 257   // CHECK-NEXT:   default: ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
 258   // CHECK-NEXT:   43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
 259   // CHECK-NEXT:   44: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
 260   // CHECK-NEXT: ]
 261     cf.switch %flag : i32, [
 262       default: ^bb2(%caseOperand0 : f32),
 263       43: ^bb3(%caseOperand1 : f32),
 264       44: ^bb4(%caseOperand2 : f32)
 265     ]
 266   ^bb2(%bb2Arg : f32):
 267     cf.br ^bb5(%bb2Arg : f32)
 268   ^bb3(%bb3Arg : f32):
 269     cf.br ^bb6(%bb3Arg : f32)
 270   ^bb4(%bb4Arg : f32):
 271     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 272
 273   // CHECK: ^[[BB5]]({{.*}}):
 274   // CHECK-NEXT: "foo.bb5Terminator"
 275   ^bb5(%bb5Arg : f32):
 276     "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
 277
 278   // CHECK: ^[[BB6]]({{.*}}):
 279   // CHECK-NEXT: "foo.bb6Terminator"
 280   ^bb6(%bb6Arg : f32):
 281     "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
 282 }
 283
 284 // CHECK-LABEL: func @switch_from_switch_with_same_value_with_match(
 285 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 286 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 287 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 288 func.func @switch_from_switch_with_same_value_with_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32) {
 289   // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
 290   "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5] : () -> ()
 291
 292   ^bb1:
 293     // CHECK: cf.switch %[[FLAG]]
 294     cf.switch %flag : i32, [
 295       default: ^bb2,
 296       42: ^bb3
 297     ]
 298
 299   ^bb2:
 300     "foo.bb2Terminator"() : () -> ()
 301   ^bb3:
 302     // prevent this block from being simplified away
 303     "foo.op"() : () -> ()
 304     // CHECK-NOT: cf.switch %[[FLAG]]
 305     // CHECK: cf.br ^[[BB5:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_1]]
 306     cf.switch %flag : i32, [
 307       default: ^bb4(%caseOperand0 : f32),
 308       42: ^bb5(%caseOperand1 : f32)
 309     ]
 310
 311   ^bb4(%bb4Arg : f32):
 312     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 313
 314   // CHECK: ^[[BB5]]({{.*}}):
 315   // CHECK-NEXT: "foo.bb5Terminator"
 316   ^bb5(%bb5Arg : f32):
 317     "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
 318 }
 319
 320 // CHECK-LABEL: func @switch_from_switch_with_same_value_no_match(
 321 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 322 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 323 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 324 // CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
 325 func.func @switch_from_switch_with_same_value_no_match(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
 326   // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
 327   "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()
 328
 329   ^bb1:
 330     // CHECK: cf.switch %[[FLAG]]
 331     cf.switch %flag : i32, [
 332       default: ^bb2,
 333       42: ^bb3
 334     ]
 335
 336   ^bb2:
 337     "foo.bb2Terminator"() : () -> ()
 338   ^bb3:
 339     "foo.op"() : () -> ()
 340     // CHECK-NOT: cf.switch %[[FLAG]]
 341     // CHECK: cf.br ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
 342     cf.switch %flag : i32, [
 343       default: ^bb4(%caseOperand0 : f32),
 344       0: ^bb5(%caseOperand1 : f32),
 345       43: ^bb6(%caseOperand2 : f32)
 346     ]
 347
 348   // CHECK: ^[[BB4]]({{.*}})
 349   // CHECK-NEXT: "foo.bb4Terminator"
 350   ^bb4(%bb4Arg : f32):
 351     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 352
 353   ^bb5(%bb5Arg : f32):
 354     "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
 355
 356   ^bb6(%bb6Arg : f32):
 357     "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
 358 }
 359
 360 // CHECK-LABEL: func @switch_from_switch_default_with_same_value(
 361 // CHECK-SAME: %[[FLAG:[a-zA-Z0-9_]+]]
 362 // CHECK-SAME: %[[CASE_OPERAND_0:[a-zA-Z0-9_]+]]
 363 // CHECK-SAME: %[[CASE_OPERAND_1:[a-zA-Z0-9_]+]]
 364 // CHECK-SAME: %[[CASE_OPERAND_2:[a-zA-Z0-9_]+]]
 365 func.func @switch_from_switch_default_with_same_value(%flag : i32, %caseOperand0 : f32, %caseOperand1 : f32, %caseOperand2 : f32) {
 366   // add predecessors for all blocks except ^bb3 to avoid other canonicalizations.
 367   "foo.pred"() [^bb1, ^bb2, ^bb4, ^bb5, ^bb6] : () -> ()
 368
 369   ^bb1:
 370     // CHECK: cf.switch %[[FLAG]]
 371     cf.switch %flag : i32, [
 372       default: ^bb3,
 373       42: ^bb2
 374     ]
 375
 376   ^bb2:
 377     "foo.bb2Terminator"() : () -> ()
 378   ^bb3:
 379     "foo.op"() : () -> ()
 380     // CHECK: cf.switch %[[FLAG]]
 381     // CHECK-NEXT: default: ^[[BB4:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_0]]
 382     // CHECK-NEXT: 43: ^[[BB6:[a-zA-Z0-9_]+]](%[[CASE_OPERAND_2]]
 383     // CHECK-NOT: 42
 384     cf.switch %flag : i32, [
 385       default: ^bb4(%caseOperand0 : f32),
 386       42: ^bb5(%caseOperand1 : f32),
 387       43: ^bb6(%caseOperand2 : f32)
 388     ]
 389
 390   // CHECK: ^[[BB4]]({{.*}}):
 391   // CHECK-NEXT: "foo.bb4Terminator"
 392   ^bb4(%bb4Arg : f32):
 393     "foo.bb4Terminator"(%bb4Arg) : (f32) -> ()
 394
 395   ^bb5(%bb5Arg : f32):
 396     "foo.bb5Terminator"(%bb5Arg) : (f32) -> ()
 397
 398   // CHECK: ^[[BB6]]({{.*}}):
 399   // CHECK-NEXT: "foo.bb6Terminator"
 400   ^bb6(%bb6Arg : f32):
 401     "foo.bb6Terminator"(%bb6Arg) : (f32) -> ()
 402 }
 403
 404 /// Test folding conditional branches that are successors of conditional
 405 /// branches with the same condition.
 406
 407 // CHECK-LABEL: func @cond_br_from_cond_br_with_same_condition
 408 func.func @cond_br_from_cond_br_with_same_condition(%cond : i1) {
 409   // CHECK:   cf.cond_br %{{.*}}, ^bb1, ^bb2
 410   // CHECK: ^bb1:
 411   // CHECK:   return
 412
 413   cf.cond_br %cond, ^bb1, ^bb2
 414
 415 ^bb1:
 416   cf.cond_br %cond, ^bb3, ^bb2
 417
 418 ^bb2:
 419   "foo.terminator"() : () -> ()
 420
 421 ^bb3:
 422   return
 423 }
 424
 425 // -----
 426
 427 // Erase assertion if condition is known to be true at compile time.
 428 // CHECK-LABEL: @assert_true
 429 func.func @assert_true() {
 430   // CHECK-NOT: cf.assert
 431   %true = arith.constant true
 432   cf.assert %true, "Computer says no"
 433   return
 434 }
 435
 436 // -----
 437
 438 // Keep assertion if condition unknown at compile time.
 439 // CHECK-LABEL: @cf.assert
 440 // CHECK-SAME:  (%[[ARG:.*]]: i1)
 441 func.func @cf.assert(%arg : i1) {
 442   // CHECK: cf.assert %[[ARG]], "Computer says no"
 443   cf.assert %arg, "Computer says no"
 444   return
 445 }
 446
 447 // -----
 448
 449 // CHECK-LABEL: @branchCondProp
 450 //       CHECK:       %[[trueval:.+]] = arith.constant true
 451 //       CHECK:       %[[falseval:.+]] = arith.constant false
 452 //       CHECK:       "test.consumer1"(%[[trueval]]) : (i1) -> ()
 453 //       CHECK:       "test.consumer2"(%[[falseval]]) : (i1) -> ()
 454 func.func @branchCondProp(%arg0: i1) {
 455   cf.cond_br %arg0, ^trueB, ^falseB
 456
 457 ^trueB:
 458   "test.consumer1"(%arg0) : (i1) -> ()
 459   cf.br ^exit
 460
 461 ^falseB:
 462   "test.consumer2"(%arg0) : (i1) -> ()
 463   cf.br ^exit
 464
 465 ^exit:
 466   return
 467 }