[LoopUnroll] Remap references in peeled iteration

[sepavloff]

Summary

If a loop contains several BBs, peeling produces code in which jump
targets point to original loop blocks rather than cloned BBs. This
patch uses the same way as loop unrolling to remap cloned blocks.

[mkuper]

lib/Transforms/Utils/LoopUnrollPeel.cpp

419 ↗

(On Diff #92781)

Something here is wrong. There's already a call to remapInstructionsInBlocks in line 443 (original) that does exactly this, so everything ought to be updated correctly. If it isn't, then something's going wrong between here and there, and I'm not sure what it is.

[sepavloff]

lib/Transforms/Utils/LoopUnrollPeel.cpp

419 ↗

(On Diff #92781)

Indeed, you are right. But the remapping is made too late, it must be done prior to dominator tree calculation, which in turn must be done prior to call to SplitBlock.

[sepavloff]

Updated patch

[mkuper]

LGTM, but please fix the test before committing.

lib/Transforms/Utils/LoopUnrollPeel.cpp

419 ↗

(On Diff #92781)

Ok, this makes more sense. So what happened here isn't that we generated bad branches, it's that we produced a malformed dominator tree?
Anyway, thanks for catching this!

test/Transforms/LoopUnroll/peel-loop.ll

99 ↗

(On Diff #92927)

Please make the test actually test its output (in particular, that we actually peeled the loop).
I think you can just run the update script on it.

[sepavloff]

Thank you for review.

lib/Transforms/Utils/LoopUnrollPeel.cpp

419 ↗

(On Diff #92781)

Yes, the problem is only malformed dominator tree. It was found in self builds with expensive checks enabled.

test/Transforms/LoopUnroll/peel-loop.ll

99 ↗

(On Diff #92927)

Done, but I had to move the test into separate file, as with options of peel-loop.ll the code changed too substantially.

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

for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     // guard here to make sure we don't end up with non-positive weights.
     if (ExitWeight < CurHeaderWeight)
       CurHeaderWeight -= ExitWeight;
     else
       CurHeaderWeight = 1;
 
     cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit,
                     NewBlocks, LoopBlocks, VMap, LVMap, DT, LI);
+
+    // Remap to use values from the current iteration instead of the
+    // previous one.
+    remapInstructionsInBlocks(NewBlocks, VMap);
+
     if (DT) {
       // Latches of the cloned loops dominate over the loop exit, so idom of the
       // latter is the first cloned loop body, as original PreHeader dominates
       // the original loop body.
       if (Iter == 0)
         DT->changeImmediateDominator(Exit, cast<BasicBlock>(LVMap[Latch]));
 #ifndef NDEBUG
       if (VerifyDomInfo)
         DT->verifyDomTree();
 #endif
     }
 
     updateBranchWeights(InsertBot, cast<BranchInst>(VMap[LatchBR]), Iter,
                         PeelCount, ExitWeight);
 
     InsertTop = InsertBot;
     InsertBot = SplitBlock(InsertBot, InsertBot->getTerminator(), DT, LI);
     InsertBot->setName(Header->getName() + ".peel.next");
 
     F->getBasicBlockList().splice(InsertTop->getIterator(),
                                   F->getBasicBlockList(),
                                   NewBlocks[0]->getIterator(), F->end());
-
-    // Remap to use values from the current iteration instead of the
-    // previous one.
-    remapInstructionsInBlocks(NewBlocks, VMap);
   }
 
   // Now adjust the phi nodes in the loop header to get their initial values
   // from the last peeled-off iteration instead of the preheader.
   for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PHI = cast<PHINode>(I);
     Value *NewVal = PHI->getIncomingValueForBlock(Latch);
     Instruction *LatchInst = dyn_cast<Instruction>(NewVal);

llvm/trunk/test/Transforms/LoopUnroll/peel-loop2.ll

+; RUN: opt -S -loop-unroll -unroll-force-peel-count=1 -verify-dom-info <%s
+
+; Check if loop composed of several BBs is peeled correctly.
+
+declare void @funcb()
+@Comma = external global i8
+define void @funca(i8* readnone %b, i8* readnone %e) {
+entry:
+  %cmp2 = icmp eq i8* %b, %e
+  br i1 %cmp2, label %for.end, label %for.body.preheader
+
+for.body.preheader:
+  br label %for.body
+
+for.body:
+  %b.addr.03 = phi i8* [ %incdec.ptr, %for.inc ], [ %b, %for.body.preheader ]
+  %0 = load i8, i8* @Comma
+  %tobool = icmp eq i8 %0, 0
+  br i1 %tobool, label %for.inc, label %if.then
+
+if.then:
+  tail call void @funcb()
+  store i8 1, i8* @Comma
+  br label %for.inc
+
+for.inc:
+  %incdec.ptr = getelementptr inbounds i8, i8* %b.addr.03, i64 1
+  %cmp = icmp eq i8* %incdec.ptr, %e
+  br i1 %cmp, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
+; CHECK_LABEL: @funca
+
+; Peeled iteration
+; CHECK: %[[REG1:[0-9]+]] = load i8, i8* @Comma
+; CHECK: %[[REG2:.*]] = icmp eq i8 %[[REG1]], 0
+; CHECK: br i1 %[[REG2]], label %{{.*}}, label %[[IFTHEN:.*]]
+; CHECK: [[IFTHEN]]:
+; CHECK: call void @funcb()
+; CHECK: store i8 1, i8* @Comma
+; CHECK: br label %[[FORINC]]
+; CHECK: [[FORINC]]:
+; CHECK: %[[REG3:.*]] = getelementptr inbounds i8, i8* %b, i64 1
+; CHECK: %[[REG4:.*]] = icmp eq i8* %[[REG3]], %e
+; CHECK: br i1 %[[REG4]]
+
+; main body
+; CHECK: %[[REG1b:.*]] = load i8, i8* @Comma
+; CHECK: %[[REG2b:.*]] = icmp eq i8 %[[REG1b]], 0
+; CHECK: br i1 %[[REG2b]], label %{{.*}}, label %[[IFTHENb:.*]]
+; CHECK: [[IFTHENb]]:
+; CHECK: call void @funcb()
+; CHECK: store i8 1, i8* @Comma
+; CHECK: br label %[[FORINCb]]
+; CHECK: [[FORINCb]]:
+; CHECK: %[[REG3b:.*]] = getelementptr inbounds i8, i8* %b, i64 1
+; CHECK: %[[REG4b:.*]] = icmp eq i8* %[[REG3b]], %e
+; CHECK: br i1 %[[REG4b]]