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( 6func.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 19func.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( 38func.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( 58func.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> 73func.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( 89func.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 107func.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( 124func.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_]+]] 148func.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_]+]] 168func.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_]+]] 193func.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_]+]] 219func.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_]+]] 247func.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_]+]] 288func.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_]+]] 325func.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_]+]] 365func.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 408func.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 429func.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) 441func.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) -> () 454func.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} 468