[RuntimeLoopUnrolling] Update DomTree correctly when exit blocks have successors

[anna]

Summary

When we runtime unroll with multiple exit blocks, we also need to update the
immediate dominators of the immediate successors of the exit blocks.

[anna]

lib/Transforms/Utils/LoopUnrollRuntime.cpp

764 ↗

(On Diff #106200)

nit: s/dominate/update
Comment should read as "we should not incorrectly *update* "

[mzolotukhin]

Don't we break critical edges before unrolling? If so, then I think this situation should never happen (and if it does, we need to understand why).

Michael

[anna]

Hi Michael,

In D35304#807108, @mzolotukhin wrote:

Don't we break critical edges before unrolling? If so, then I think this situation should never happen (and if it does, we need to understand why).

Michael

As part of support for runtime unrolling in presence of multiple exit blocks (functionality was added in rL306846), when we have multiple exits in a loop, we generate edges from both the main loop and the remainder loop to the main loop's exit block (this is the structure after cloneLoopBlocks). This was to avoid cloning of these multiple exit blocks. After this is done, there are three cleanup steps done:

1. update the phi nodes in the exit blocks.
2. update the immediate dominator of these exit blocks
3. preserve the hasDedicatedExits property for both main and remainder loop: see lines 848-856.

Does this answer your question? Thanks.

[mzolotukhin]

Does this answer your question?

Yes, thank you for the detailed answer. The patch looks good to me then (please also find some minor comments/suggestions inline).

Thanks,
Michael

lib/Transforms/Utils/LoopUnrollRuntime.cpp

768–773 ↗

(On Diff #106200)

Are there any other situations where we might want to check this property? It might make sense to check it even before we checked if DT is available (under #ifdef EXPENSIVE_CHECKS of course).

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

8–10 ↗

(On Diff #106200)

This wording seems a bit strange. Is one of them a leftover from some editing?

23 ↗

(On Diff #106200)

I'd prefer to simplify br i1 true to unconditional branch or introduce a function parameter if you really need this branch to be conditional. Same applies to undefs used later in other places.

25 ↗

(On Diff #106200)

Redundant empty line.

83–85 ↗

(On Diff #106200)

Could you please rearrange blocks, so that they stay in a (topo)logical order (like entry, preheader, header, latch, exit1, exit2, mergedexit)?

131 ↗

(On Diff #106200)

Do we really need the shift in this test?

[anna]

lib/Transforms/Utils/LoopUnrollRuntime.cpp

768–773 ↗

(On Diff #106200)

While running this through make-check, I realized the assertion should be the inverse: if the LHS fails, then we trip on the assert. Just like the assert on line 775 :)

assert(!(any_of(OtherExits,
                     [SuccBB](BasicBlock *EB) { return EB == SuccBB; }) ||
                SuccBB == LatchExit) &&
              "Breaks the definition of dedicated exits!");

I couldn't think of any other situations where we need to check this property about successors of the exit blocks (intent is clearer within this context). Still this is a valid property outside of DT, so moved it outside of DT.

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

8–10 ↗

(On Diff #106200)

yes, thanks!

83–85 ↗

(On Diff #106200)

yes, done for all test cases.

131 ↗

(On Diff #106200)

It can be any other operation. It's used to show the phi in the latchexit.

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

 //
 // The current strategy generates an if-then-else sequence prior to the
 // unrolled loop to execute the 'left over' iterations before or after the
 // unrolled loop.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/LoopPass.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Metadata.h"
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, since we have new edges into OtherExits.
+  // OtherExits and their immediate successors, since we have new edges into
+  // OtherExits.
+  SmallSet<BasicBlock*, 8> ImmediateSuccessorsOfExitBlocks;
   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;
      PHINode *Phi = cast<PHINode>(&II);
      unsigned oldNumOperands = Phi->getNumIncomingValues();
      // Add the incoming values from the remainder code to the end of the phi
      // node.
      for (unsigned i =0; i < oldNumOperands; i++){
        Value *newVal = VMap[Phi->getIncomingValue(i)];
        // newVal can be a constant or derived from values outside the loop, and
        // hence need not have a VMap value.
        if (!newVal)
          newVal = Phi->getIncomingValue(i);
        Phi->addIncoming(newVal,
                            cast<BasicBlock>(VMap[Phi->getIncomingBlock(i)]));
      }
    }
+#ifdef EXPENSIVE_CHECKS
+    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)
+   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);
+         }
+       }
+    }
   }
 
   // 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

+; RUN: opt < %s -loop-unroll -unroll-runtime=true -unroll-runtime-epilog=false -unroll-runtime-multi-exit=true -unroll-count=4  -verify-dom-info -S | FileCheck %s
+
+; REQUIRES: asserts
+; The tests below are for verifying dom tree after runtime unrolling
+; with multiple exit/exiting blocks.
+
+; We explicitly set the unroll count so that expensiveTripCount computation is allowed.
+
+; mergedexit block has edges from loop exit blocks.
+define i64 @test1() {
+; CHECK-LABEL: test1(
+; CHECK-LABEL: headerexit:
+; CHECK-NEXT:    %addphi = phi i64 [ %add.iv, %header ], [ %add.iv.1, %header.1 ], [ %add.iv.2, %header.2 ], [ %add.iv.3, %header.3 ]
+; CHECK-NEXT:    br label %mergedexit
+; CHECK-LABEL: latchexit:
+; CHECK-NEXT:    %shftphi = phi i64 [ %shft, %latch ], [ %shft.1, %latch.1 ], [ %shft.2, %latch.2 ], [ %shft.3, %latch.3 ]
+; CHECK-NEXT:    br label %mergedexit
+; CHECK-LABEL: mergedexit:
+; CHECK-NEXT:    %retval = phi i64 [ %addphi, %headerexit ], [ %shftphi, %latchexit ]
+; CHECK-NEXT:    ret i64 %retval
+entry:
+  br label %preheader
+
+preheader:                                              ; preds = %bb
+  %trip = zext i32 undef to i64
+  br label %header
+
+header:                                              ; preds = %latch, %preheader
+  %iv = phi i64 [ 2, %preheader ], [ %add.iv, %latch ]
+  %add.iv = add nuw nsw i64 %iv, 2
+  %cmp1 = icmp ult i64 %add.iv, %trip
+  br i1 %cmp1, label %latch, label %headerexit
+
+latch:                                             ; preds = %header
+  %shft = ashr i64 %add.iv, 1
+  %cmp2 = icmp ult i64 %shft, %trip
+  br i1 %cmp2, label %header, label %latchexit
+
+headerexit:                                              ; preds = %header
+  %addphi = phi i64 [ %add.iv, %header ]
+  br label %mergedexit
+
+latchexit:                                              ; preds = %latch
+ %shftphi = phi i64 [ %shft, %latch ]
+  br label %mergedexit
+
+mergedexit:                                              ; preds = %latchexit, %headerexit
+  %retval = phi i64 [ %addphi, %headerexit ], [ %shftphi, %latchexit ]
+  ret i64 %retval
+}
+
+; mergedexit has edges from loop exit blocks and a block outside the loop.
+define  void @test2(i1 %cond, i32 %n) {
+; CHECK-LABEL: header.1:
+; CHECK-NEXT:    %add.iv.1 = add nuw nsw i64 %add.iv, 2
+; CHECK:         br i1 %cmp1.1, label %latch.1, label %headerexit
+; CHECK-LABEL: latch.3:
+; CHECK:         %cmp2.3 = icmp ult i64 %shft.3, %trip
+; CHECK-NEXT:    br i1 %cmp2.3, label %header, label %latchexit, !llvm.loop
+entry:
+  br i1 %cond, label %preheader, label %mergedexit
+
+preheader:                                              ; preds = %entry
+  %trip = zext i32 %n to i64
+  br label %header
+
+header:                                              ; preds = %latch, %preheader
+  %iv = phi i64 [ 2, %preheader ], [ %add.iv, %latch ]
+  %add.iv = add nuw nsw i64 %iv, 2
+  %cmp1 = icmp ult i64 %add.iv, %trip
+  br i1 %cmp1, label %latch, label %headerexit
+
+latch:                                             ; preds = %header
+  %shft = ashr i64 %add.iv, 1
+  %cmp2 = icmp ult i64 %shft, %trip
+  br i1 %cmp2, label %header, label %latchexit
+
+headerexit:                                              ; preds = %header
+  br label %mergedexit
+
+latchexit:                                              ; preds = %latch
+  br label %mergedexit
+
+mergedexit:                                              ; preds = %latchexit, %headerexit, %entry
+  ret void
+}
+
+
+; exitsucc is from loop exit block only.
+define i64 @test3(i32 %n) {
+; CHECK-LABEL: test3(
+; CHECK-LABEL:  headerexit:
+; CHECK-NEXT:     br label %exitsucc
+; CHECK-LABEL:  latchexit:
+; CHECK-NEXT:     %shftphi = phi i64 [ %shft, %latch ], [ %shft.1, %latch.1 ], [ %shft.2, %latch.2 ], [ %shft.3, %latch.3 ]
+; CHECK-NEXT:     ret i64 %shftphi
+; CHECK-LABEL:  exitsucc:
+; CHECK-NEXT:     ret i64 96
+entry:
+  br label %preheader
+
+preheader:                                              ; preds = %bb
+  %trip = zext i32 %n to i64
+  br label %header
+
+header:                                              ; preds = %latch, %preheader
+  %iv = phi i64 [ 2, %preheader ], [ %add.iv, %latch ]
+  %add.iv = add nuw nsw i64 %iv, 2
+  %cmp1 = icmp ult i64 %add.iv, %trip
+  br i1 %cmp1, label %latch, label %headerexit
+
+latch:                                             ; preds = %header
+  %shft = ashr i64 %add.iv, 1
+  %cmp2 = icmp ult i64 %shft, %trip
+  br i1 %cmp2, label %header, label %latchexit
+
+headerexit:                                              ; preds = %header
+  br label %exitsucc
+
+latchexit:                                              ; preds = %latch
+  %shftphi = phi i64 [ %shft, %latch ]
+  ret i64 %shftphi
+
+exitsucc:                                              ; preds = %headerexit
+  ret i64 96
+}