Skip to content

Commit 631b89c

Browse files
benwu25benwu25
authored
[MLIR][CF] Avoid collapsing blocks which participate in cycles (#160783)
Previously, collapseBranch did not return failure for successor blocks which were part of a cycle. mlir-opt --canonicalize would run indefinitely for any N-block cycle which is kicked off with an unconditional jump. The simplifyPassThroughBr transform would continue alternating which block was targeted in ^bb0, resulting in an infinite loop. collapseBranch will not result in any useful transformation on blocks which participate in cycles, since the block is aliased by a different block. To avoid this, we can check for cycles in collapseBranch and abort when one is detected. Simplification of the cycle is left for other transforms. Fixes #159743.
1 parent 2e19666 commit 631b89c

File tree

3 files changed

+158
-2
lines changed

3 files changed

+158
-2
lines changed

‎flang/test/Lower/OpenMP/infinite-loop-in-construct.f90‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,10 @@
88
! CHECK: cf.cond_br %{{[0-9]+}}, ^bb1, ^bb2
99
! CHECK-NEXT: ^bb1: // pred: ^bb0
1010
! CHECK: cf.br ^bb2
11-
! CHECK-NEXT: ^bb2: // 3 preds: ^bb0, ^bb1, ^bb2
12-
! CHECK-NEXT: cf.br ^bb2
11+
! CHECK-NEXT: ^bb2: // 2 preds: ^bb0, ^bb1
12+
! CHECK: cf.br ^bb3
13+
! CHECK-NEXT: ^bb3: // 2 preds: ^bb2, ^bb3
14+
! CHECK: cf.br ^bb3
1315
! CHECK-NEXT: }
1416

1517
subroutine sb(ninter, numnod)

‎mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,16 @@ static LogicalResult collapseBranch(Block *&successor,
122122
Block *successorDest = successorBranch.getDest();
123123
if (successorDest == successor)
124124
return failure();
125+
// Don't try to collapse branches which participate in a cycle.
126+
BranchOp nextBranch = dyn_cast<BranchOp>(successorDest->getTerminator());
127+
llvm::DenseSet<Block *> visited{successor, successorDest};
128+
while (nextBranch) {
129+
Block *nextBranchDest = nextBranch.getDest();
130+
if (visited.contains(nextBranchDest))
131+
return failure();
132+
visited.insert(nextBranchDest);
133+
nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
134+
}
125135

126136
// Update the operands to the successor. If the branch parent has no
127137
// arguments, we can use the branch operands directly.

‎mlir/test/Dialect/ControlFlow/canonicalize.mlir‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -490,3 +490,147 @@ func.func @branchCondProp(%arg0: i1) {
490490
^exit:
491491
return
492492
}
493+
494+
// -----
495+
496+
/// Test that control-flow cycles are not simplified infinitely.
497+
498+
// CHECK-LABEL: @cycle_2_blocks
499+
// CHECK: cf.br ^bb1
500+
// CHECK: ^bb1:
501+
// CHECK: cf.br ^bb1
502+
func.func @cycle_2_blocks() {
503+
cf.br ^bb1
504+
^bb1:
505+
cf.br ^bb2
506+
^bb2:
507+
cf.br ^bb1
508+
}
509+
510+
// CHECK-LABEL: @no_cycle_2_blocks
511+
// CHECK: %[[VAL_0:.*]] = arith.constant 1 : i32
512+
// CHECK: return %[[VAL_0]] : i32
513+
func.func @no_cycle_2_blocks() -> i32 {
514+
cf.br ^bb1
515+
^bb1:
516+
cf.br ^bb2
517+
^bb2:
518+
cf.br ^bb3
519+
^bb3:
520+
%ret = arith.constant 1 : i32
521+
return %ret : i32
522+
}
523+
524+
// CHECK-LABEL: @cycle_4_blocks
525+
// CHECK: cf.br ^bb1
526+
// CHECK: ^bb1:
527+
// CHECK: cf.br ^bb1
528+
func.func @cycle_4_blocks() {
529+
cf.br ^bb1
530+
^bb1:
531+
cf.br ^bb2
532+
^bb2:
533+
cf.br ^bb3
534+
^bb3:
535+
cf.br ^bb4
536+
^bb4:
537+
cf.br ^bb1
538+
}
539+
540+
// CHECK-LABEL: @no_cycle_4_blocks
541+
// CHECK: %[[VAL_0:.*]] = arith.constant 1 : i32
542+
// CHECK: return %[[VAL_0]] : i32
543+
func.func @no_cycle_4_blocks() -> i32 {
544+
cf.br ^bb1
545+
^bb1:
546+
cf.br ^bb2
547+
^bb2:
548+
cf.br ^bb3
549+
^bb3:
550+
cf.br ^bb4
551+
^bb4:
552+
cf.br ^bb5
553+
^bb5:
554+
%ret = arith.constant 1 : i32
555+
return %ret : i32
556+
}
557+
558+
// CHECK-LABEL: @delayed_3_cycle
559+
// CHECK: cf.br ^bb1
560+
// CHECK: ^bb1:
561+
// CHECK: cf.br ^bb1
562+
func.func @delayed_3_cycle() {
563+
cf.br ^bb1
564+
^bb1:
565+
cf.br ^bb2
566+
^bb2:
567+
cf.br ^bb3
568+
^bb3:
569+
cf.br ^bb4
570+
^bb4:
571+
cf.br ^bb5
572+
^bb5:
573+
cf.br ^bb3
574+
}
575+
576+
// CHECK-LABEL: @cycle_1_block
577+
// CHECK: cf.br ^bb1
578+
// CHECK: ^bb1:
579+
// CHECK: cf.br ^bb1
580+
func.func @cycle_1_block() {
581+
cf.br ^bb1
582+
^bb1:
583+
cf.br ^bb2
584+
^bb2:
585+
cf.br ^bb2
586+
}
587+
588+
// CHECK-LABEL: @unsimplified_cycle_1
589+
// CHECK-SAME: %[[ARG0:.*]]: i1) {
590+
// CHECK: cf.cond_br %[[ARG0]], ^bb1, ^bb2
591+
// CHECK: ^bb1:
592+
// CHECK: cf.br ^bb2
593+
// CHECK: ^bb2:
594+
// CHECK: cf.br ^bb3
595+
// CHECK: ^bb3:
596+
// CHECK: cf.br ^bb3
597+
func.func @unsimplified_cycle_1(%c : i1) {
598+
cf.cond_br %c, ^bb1, ^bb2
599+
^bb1:
600+
cf.br ^bb2
601+
^bb2:
602+
cf.br ^bb3
603+
^bb3:
604+
cf.br ^bb4
605+
^bb4:
606+
cf.br ^bb3
607+
}
608+
609+
// Make sure we terminate when other cf passes can't help us.
610+
611+
// CHECK-LABEL: @unsimplified_cycle_2
612+
// CHECK-SAME: %[[ARG0:.*]]: i1) {
613+
// CHECK: cf.cond_br %[[ARG0]], ^bb1, ^bb3
614+
// CHECK: ^bb1:
615+
// CHECK: cf.br ^bb2 {A}
616+
// CHECK: ^bb2:
617+
// CHECK: cf.br ^bb2 {E}
618+
// CHECK: ^bb3:
619+
// CHECK: cf.br ^bb1
620+
func.func @unsimplified_cycle_2(%c : i1) {
621+
cf.cond_br %c, ^bb6, ^bb7
622+
^bb6:
623+
cf.br ^bb5 {F}
624+
^bb5:
625+
cf.br ^bb1 {A}
626+
^bb1:
627+
cf.br ^bb2 {B}
628+
^bb2:
629+
cf.br ^bb3 {C}
630+
^bb3:
631+
cf.br ^bb4 {D}
632+
^bb4:
633+
cf.br ^bb1 {E}
634+
^bb7:
635+
cf.br ^bb6
636+
}

0 commit comments

Comments
 (0)