[UnrollRuntime] Fix domTree failure in multiexit unrolling

[anna]

Summary

This fixes the IDom for an exit block when runtime unrolling under multiexit/exiting case.
We initially had a restrictive check that the IDom is only updated when
it is the header of the loop.
However, we also need to update the IDom to the correct one when the
IDom is any block within the original loop. See added test case (which
fails dom tree verification without the patch).

[anna]

Context and history for runtime unrolling: https://reviews.llvm.org/D35304

[brzycki]

LGTM. @kuhar may wish to comment.

[kuhar]

@anna Can this be rewritten to use the DomTreeUpdater utility? We have large number of examples in the codebase already, so the question is whether this affects the performance in this case. Historically, we had tons of bugs in manual DomTree updates and I'd rather we replaces as much as possible using the incremental updater.

[anna]

In D56284#1345766, @kuhar wrote:

@anna Can this be rewritten to use the DomTreeUpdater utility?

@kuhar, I have not looked at the DTU uttility yet, but I'd think it can be rewritten (unless there are some limitations for the DomTreeUpdater utility). However, that would be a large enough change and I think it's better to do it as a separate change at a later point rather than as part of fixing this bug. Also, there are enough DT updates through out this code that it will take a while to get through the whole process of porting over for runtime unrolling.

[kuhar]

In D56284#1345811, @anna wrote:

In D56284#1345766, @kuhar wrote:

@anna Can this be rewritten to use the DomTreeUpdater utility?

@kuhar, I have not looked at the DTU uttility yet, but I'd think it can be rewritten (unless there are some limitations for the DomTreeUpdater utility). However, that would be a large enough change and I think it's better to do it as a separate change at a later point rather than as part of fixing this bug. Also, there are enough DT updates through out this code that it will take a while to get through the whole process of porting over for runtime unrolling.

I see, thank you for the explanation. I think think there's nothing stopping it from being ported, this will be the correct long-term fix. I skimmed through the code and it shouldn't take more than a day. But as you mentioned, it can happen in a separate patch.

[anna]

The IDom can be any block within the loop (as mentioned in the comment).
So, we need to check for IDom *contained* in the loop.
Added test case shows where the IDom is within the inner loop of the original loop.
Both test cases pass now with the fix.

[anna]

the same bug can happen with:

1. latch exit
2. non-immediate successors of latchexit/otherexit.

I have simplified test cases for each of these and we need a more general fix. Working on this.

[kuhar]

In D56284#1348681, @anna wrote:

the same bug can happen with:

1. latch exit
2. non-immediate successors of latchexit/otherexit.

I have simplified test cases for each of these and we need a more general fix. Working on this.

... these things are extremely bug-prone; I really do suggest using the automatic updater instead of trying to deal with all tricky corner-cases here :)

[anna]

In D56284#1348693, @kuhar wrote:

In D56284#1348681, @anna wrote:

the same bug can happen with:

1. latch exit
2. non-immediate successors of latchexit/otherexit.

I have simplified test cases for each of these and we need a more general fix. Working on this.

... these things are extremely bug-prone; I really do suggest using the automatic updater instead of trying to deal with all tricky corner-cases here :)

I've got a naive question about the DTU: Does the automatic updater handle the ImmediateDominator "automatically" for the remaining nodes in the dom tree once identify that one of the nodes in the DT has a change in IDom?

This will drastically help with #2 because what I have as a local fix using the old DT is something like this:

+   if (DT) {
+    // Check the dom children of each block in loop and if it is outside the
+    // current loop, update it to the preheader.
+    for (auto *BB: L->blocks()) {
+      auto *DomNodeBB = DT->getNode(BB);
+      for (auto *DomChild: DomNodeBB->getChildren()) {
+        if (!L->contains(LI->getLoopFor(DomChild->getBlock()))) {
+          DT->changeImmediateDominator(DomChild, DT->getNode(PreHeader));
+        }    
+      }    
+    }
+   }

This is a more general fix using the old DT to handle the non-immediate successors of an exit block in the loop.

[kuhar]

In D56284#1348747, @anna wrote:

I've got a naive question about the DTU: Does the automatic updater handle the ImmediateDominator "automatically" for the remaining nodes in the dom tree once identify that one of the nodes in the DT has a change in IDom?

Exactly. Given a set of CFG updates, if figures out how to change the DomTree and PostDomTree such that they match the CFG.

This will drastically help with #2 because what I have as a local fix using the old DT is something like this:

+   if (DT) {
+    // Check the dom children of each block in loop and if it is outside the
+    // current loop, update it to the preheader.
+    for (auto *BB: L->blocks()) {
+      auto *DomNodeBB = DT->getNode(BB);
+      for (auto *DomChild: DomNodeBB->getChildren()) {
+        if (!L->contains(LI->getLoopFor(DomChild->getBlock()))) {
+          DT->changeImmediateDominator(DomChild, DT->getNode(PreHeader));
+        }    
+      }    
+    }
+   }

This is a more general fix using the old DT to handle the non-immediate successors of an exit block in the loop.

Doesn't this loop have to be run until you reach a fixpoint?

[anna]

In D56284#1348783, @kuhar wrote:

In D56284#1348747, @anna wrote:

I've got a naive question about the DTU: Does the automatic updater handle the ImmediateDominator "automatically" for the remaining nodes in the dom tree once identify that one of the nodes in the DT has a change in IDom?

Exactly. Given a set of CFG updates, if figures out how to change the DomTree and PostDomTree such that they match the CFG.

This will drastically help with #2 because what I have as a local fix using the old DT is something like this:

+   if (DT) {
+    // Check the dom children of each block in loop and if it is outside the
+    // current loop, update it to the preheader.
+    for (auto *BB: L->blocks()) {
+      auto *DomNodeBB = DT->getNode(BB);
+      for (auto *DomChild: DomNodeBB->getChildren()) {
+        if (!L->contains(LI->getLoopFor(DomChild->getBlock()))) {
+          DT->changeImmediateDominator(DomChild, DT->getNode(PreHeader));
+        }    
+      }    
+    }
+   }

This is a more general fix using the old DT to handle the non-immediate successors of an exit block in the loop.

Doesn't this loop have to be run until you reach a fixpoint?

yes, that's right - I'd missed that. Now, It's a principled fix which handles all cases of nearest common dominator of the IDoms of multiple exiting blocks, and these other cases I've seen above.

If any more DT bugs shows up in our fuzzer run, I'm just going to scrap this and go with the DTU itself instead of the other way around - landing this and then changing to DTU :)

[anna]

we need to update domInfo for exits and those reachable from the exits.

[anna]

This has passed all internal fuzzer runs - the DT verification failure tests have been added as test cases here after reduction.

[brzycki]

@anna thank you for your patience and perseverance with this fix and welcome to the hairy world of manual dominator updates. :)

I too would like to see this moved to DTU in a future update but for now I feel it's sufficient that it's passed your fuzzer tests and is much better than the current state of DT verification failure.

LGTM.

[anna]

thanks for reviewing the change Brian! Will land it soon.

[brzycki]

I am accepting this revision.

llvm/trunk/lib/Transforms/Utils/LoopUnrollRuntime.cpp

bool llvm::UnrollRuntimeLoopRemainder(Loop *L, unsigned Count,
   F->getBasicBlockList().splice(InsertBot->getIterator(),
                                 F->getBasicBlockList(),
                                 NewBlocks[0]->getIterator(),
                                 F->end());
 
   // Now the loop blocks are cloned and the other exiting blocks from the
   // remainder are connected to the original Loop's exit blocks. The remaining
   // work is to update the phi nodes in the original loop, and take in the
-  // values from the cloned region. Also update the dominator info for
-  // OtherExits and their immediate successors, since we have new edges into
-  // OtherExits.
-  SmallPtrSet<BasicBlock*, 8> ImmediateSuccessorsOfExitBlocks;
+  // values from the cloned region.
   for (auto *BB : OtherExits) {
    for (auto &II : *BB) {
 
      // Given we preserve LCSSA form, we know that the values used outside the
      // loop will be used through these phi nodes at the exit blocks that are
      // transformed below.
      if (!isa<PHINode>(II))
        break;

 #if defined(EXPENSIVE_CHECKS) && !defined(NDEBUG)
     for (BasicBlock *SuccBB : successors(BB)) {
       assert(!(any_of(OtherExits,
                       [SuccBB](BasicBlock *EB) { return EB == SuccBB; }) ||
                SuccBB == LatchExit) &&
              "Breaks the definition of dedicated exits!");
     }
 #endif
-   // Update the dominator info because the immediate dominator is no longer the
-   // header of the original Loop. BB has edges both from L and remainder code.
-   // Since the preheader determines which loop is run (L or directly jump to
-   // the remainder code), we set the immediate dominator as the preheader.
-   if (DT) {
-     DT->changeImmediateDominator(BB, PreHeader);
-     // Also update the IDom for immediate successors of BB.  If the current
-     // IDom is the header, update the IDom to be the preheader because that is
-     // the nearest common dominator of all predecessors of SuccBB.  We need to
-     // check for IDom being the header because successors of exit blocks can
-     // have edges from outside the loop, and we should not incorrectly update
-     // the IDom in that case.
-     for (BasicBlock *SuccBB: successors(BB))
-       if (ImmediateSuccessorsOfExitBlocks.insert(SuccBB).second) {
-         if (DT->getNode(SuccBB)->getIDom()->getBlock() == Header) {
-           assert(!SuccBB->getSinglePredecessor() &&
-                  "BB should be the IDom then!");
-           DT->changeImmediateDominator(SuccBB, PreHeader);
-         }
-       }
+
+  // Update the immediate dominator of the exit blocks and blocks that are
+  // reachable from the exit blocks. This is needed because we now have paths
+  // from both the original loop and the remainder code reaching the exit
+  // blocks. While the IDom of these exit blocks were from the original loop,
+  // now the IDom is the preheader (which decides whether the original loop or
+  // remainder code should run).
+  if (DT && !L->getExitingBlock()) {
+    SmallVector<BasicBlock *, 16> ChildrenToUpdate;
+    // NB! We have to examine the dom children of all loop blocks, not just
+    // those which are the IDom of the exit blocks. This is because blocks
+    // reachable from the exit blocks can have their IDom as the nearest common
+    // dominator of the exit blocks.
+    for (auto *BB : L->blocks()) {
+      auto *DomNodeBB = DT->getNode(BB);
+      for (auto *DomChild : DomNodeBB->getChildren()) {
+        auto *DomChildBB = DomChild->getBlock();
+        if (!L->contains(LI->getLoopFor(DomChildBB)))
+          ChildrenToUpdate.push_back(DomChildBB);
     }
+    for (auto *BB : ChildrenToUpdate)
+      DT->changeImmediateDominator(BB, PreHeader);
   }
 
   // Loop structure should be the following:
   //  Epilog             Prolog
   //
   // PreHeader         PreHeader
   // NewPreHeader      PrologPreHeader
   //   Header            PrologHeader

llvm/trunk/test/Transforms/LoopUnroll/runtime-loop-multiexit-dom-verify.ll

 
 latchexit:                                              ; preds = %latch
   %shftphi = phi i64 [ %shft, %latch ]
   ret i64 %shftphi
 
 exitsucc:                                              ; preds = %headerexit
   ret i64 96
 }
+
+; exit block (%default) has an exiting block and another exit block as predecessors.
+define void @test4(i16 %c3) {
+; CHECK-LABEL: test4
+
+; CHECK-LABEL: exiting.prol:
+; CHECK-NEXT:   switch i16 %c3, label %default.loopexit.loopexit1 [
+
+; CHECK-LABEL: exiting:
+; CHECK-NEXT:   switch i16 %c3, label %default.loopexit.loopexit [
+
+; CHECK-LABEL: default.loopexit.loopexit:
+; CHECK-NEXT:   br label %default.loopexit
+
+; CHECK-LABEL: default.loopexit.loopexit1:
+; CHECK-NEXT:   br label %default.loopexit
+
+; CHECK-LABEL: default.loopexit:
+; CHECK-NEXT:   br label %default
+preheader:
+  %c1 = zext i32 undef to i64
+  br label %header
+
+header:                                       ; preds = %latch, %preheader
+  %indvars.iv = phi i64 [ 0, %preheader ], [ %indvars.iv.next, %latch ]
+  br label %exiting
+
+exiting:                                           ; preds = %header
+  switch i16 %c3, label %default [
+    i16 45, label %otherexit
+    i16 95, label %latch
+  ]
+
+latch:                                          ; preds = %exiting
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %c2 = icmp ult i64 %indvars.iv.next, %c1
+  br i1 %c2, label %header, label %latchexit
+
+latchexit:                                          ; preds = %latch
+  ret void
+
+default:                                          ; preds = %otherexit, %exiting
+  ret void
+
+otherexit:                                           ; preds = %exiting
+  br label %default
+}
+
+; exit block (%exitB) has an exiting block and another exit block as predecessors.
+; exiting block comes from inner loop.
+define void @test5() {
+; CHECK-LABEL: test5
+; CHECK-LABEL: bb1:
+; CHECK-NEXT:   br i1 false, label %outerH.prol.preheader, label %outerH.prol.loopexit
+
+; CHECK-LABEL: outerH.prol.preheader:
+; CHECK-NEXT:   br label %outerH.prol
+
+; CHECK-LABEL: outerH.prol:
+; CHECK-NEXT:   %tmp4.prol = phi i32 [ %tmp6.prol, %outerLatch.prol ], [ undef, %outerH.prol.preheader ]
+; CHECK-NEXT:   %prol.iter = phi i32 [ 0, %outerH.prol.preheader ], [ %prol.iter.sub, %outerLatch.prol ]
+; CHECK-NEXT:   br label %innerH.prol
+bb:
+  %tmp = icmp sgt i32 undef, 79
+  br i1 %tmp, label %outerLatchExit, label %bb1
+
+bb1:                                              ; preds = %bb
+  br label %outerH
+
+outerH:                                              ; preds = %outerLatch, %bb1
+  %tmp4 = phi i32 [ %tmp6, %outerLatch ], [ undef, %bb1 ]
+  br label %innerH
+
+innerH:                                              ; preds = %innerLatch, %outerH
+  br i1 undef, label %innerexiting, label %otherexitB
+
+innerexiting:                                             ; preds = %innerH
+  br i1 undef, label %innerLatch, label %exitB
+
+innerLatch:                                             ; preds = %innerexiting
+  %tmp13 = fcmp olt double undef, 2.000000e+00
+  br i1 %tmp13, label %innerH, label %outerLatch
+
+outerLatch:                                              ; preds = %innerLatch
+  %tmp6 = add i32 %tmp4, 1
+  %tmp7 = icmp sgt i32 %tmp6, 79
+  br i1 %tmp7, label %outerLatchExit, label %outerH
+
+outerLatchExit:                                              ; preds = %outerLatch, %bb
+  ret void
+
+exitB:                                             ; preds = %innerexiting, %otherexitB
+  ret void
+
+otherexitB:                                              ; preds = %innerH
+  br label %exitB
+
+}
+
+; Blocks reachable from exits (not_zero44) have the IDom as the block within the loop (Header).
+; Update the IDom to the preheader.
+define void @test6() {
+; CHECK-LABEL: test6
+; CHECK-LABEL: header.prol.preheader:
+; CHECK-NEXT:    br label %header.prol
+
+; CHECK-LABEL: header.prol:
+; CHECK-NEXT:    %indvars.iv.prol = phi i64 [ undef, %header.prol.preheader ], [ %indvars.iv.next.prol, %latch.prol ]
+; CHECK-NEXT:    %prol.iter = phi i64 [ 1, %header.prol.preheader ], [ %prol.iter.sub, %latch.prol ]
+; CHECK-NEXT:    br i1 false, label %latch.prol, label %otherexit.loopexit1
+
+; CHECK-LABEL: header.prol.loopexit.unr-lcssa:
+; CHECK-NEXT:    %indvars.iv.unr.ph = phi i64 [ %indvars.iv.next.prol, %latch.prol ]
+; CHECK-NEXT:    br label %header.prol.loopexit
+
+; CHECK-LABEL: header.prol.loopexit:
+; CHECK-NEXT:    %indvars.iv.unr = phi i64 [ undef, %entry ], [ %indvars.iv.unr.ph, %header.prol.loopexit.unr-lcssa ]
+; CHECK-NEXT:    br i1 true, label %latchexit, label %entry.new
+
+; CHECK-LABEL: entry.new:
+; CHECK-NEXT:    br label %header
+entry:
+  br label %header
+
+header:                                          ; preds = %latch, %entry
+  %indvars.iv = phi i64 [ undef, %entry ], [ %indvars.iv.next, %latch ]
+  br i1 undef, label %latch, label %otherexit
+
+latch:                                         ; preds = %header
+  %indvars.iv.next = add nsw i64 %indvars.iv, 2
+  %0 = icmp slt i64 %indvars.iv.next, 616
+  br i1 %0, label %header, label %latchexit
+
+latchexit:                                          ; preds = %latch
+  br label %latchexitsucc
+
+otherexit:                                 ; preds = %header
+  br label %otherexitsucc
+
+otherexitsucc:                                          ; preds = %otherexit
+  br label %not_zero44
+
+not_zero44:                                       ; preds = %latchexitsucc, %otherexitsucc
+  unreachable
+
+latchexitsucc:                                      ; preds = %latchexit
+  br label %not_zero44
+}
+