Index: lib/Transforms/Utils/LoopUnrollRuntime.cpp =================================================================== --- lib/Transforms/Utils/LoopUnrollRuntime.cpp +++ lib/Transforms/Utils/LoopUnrollRuntime.cpp @@ -36,6 +36,7 @@ #include "llvm/Transforms/Scalar.h" #include "llvm/Transforms/Utils/BasicBlockUtils.h" #include "llvm/Transforms/Utils/Cloning.h" +#include "llvm/Transforms/Utils/LoopUtils.h" #include "llvm/Transforms/Utils/UnrollLoop.h" #include @@ -45,6 +46,10 @@ STATISTIC(NumRuntimeUnrolled, "Number of loops unrolled with run-time trip counts"); +static cl::opt UnrollRuntimeMultiExit( + "unroll-runtime-multi-exit", cl::init(false), cl::Hidden, + cl::desc("Allow runtime unrolling for loops with multiple exits, when " + "epilog is generated")); /// Connect the unrolling prolog code to the original loop. /// The unrolling prolog code contains code to execute the @@ -286,14 +291,12 @@ /// If loop structure is cloned InsertTop should be new preheader, InsertBot /// new loop exit. /// -static void CloneLoopBlocks(Loop *L, Value *NewIter, - const bool CreateRemainderLoop, - const bool UseEpilogRemainder, - BasicBlock *InsertTop, BasicBlock *InsertBot, - BasicBlock *Preheader, - std::vector &NewBlocks, - LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap, - DominatorTree *DT, LoopInfo *LI) { +static Loop * +CloneLoopBlocks(Loop *L, Value *NewIter, const bool CreateRemainderLoop, + const bool UseEpilogRemainder, BasicBlock *InsertTop, + BasicBlock *InsertBot, BasicBlock *Preheader, + std::vector &NewBlocks, LoopBlocksDFS &LoopBlocks, + ValueToValueMapTy &VMap, DominatorTree *DT, LoopInfo *LI) { StringRef suffix = UseEpilogRemainder ? "epil" : "prol"; BasicBlock *Header = L->getHeader(); BasicBlock *Latch = L->getLoopLatch(); @@ -418,7 +421,10 @@ // Set operand 0 to refer to the loop id itself. NewLoopID->replaceOperandWith(0, NewLoopID); NewLoop->setLoopID(NewLoopID); + return NewLoop; } + else + return nullptr; } /// Insert code in the prolog/epilog code when unrolling a loop with a @@ -465,29 +471,48 @@ LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, bool PreserveLCSSA) { // for now, only unroll loops that contain a single exit - if (!L->getExitingBlock()) + if (!UnrollRuntimeMultiExit && !L->getExitingBlock()) return false; - // Make sure the loop is in canonical form, and there is a single - // exit block only. + // Make sure the loop is in canonical form. if (!L->isLoopSimplifyForm()) return false; // Guaranteed by LoopSimplifyForm. BasicBlock *Latch = L->getLoopLatch(); + BasicBlock *Header = L->getHeader(); BasicBlock *LatchExit = L->getUniqueExitBlock(); // successor out of loop - if (!LatchExit) + if (!LatchExit && !UnrollRuntimeMultiExit) return false; + // These are exit blocks other than the target of the latch exiting block. + SmallVector OtherExits; + BranchInst *LatchBR = cast(Latch->getTerminator()); + unsigned int ExitIndex = LatchBR->getSuccessor(0) == Header ? 1 : 0; // Cloning the loop basic blocks (`CloneLoopBlocks`) requires that one of the - // targets of the Latch be the single exit block out of the loop. This needs + // targets of the Latch be an exit block out of the loop. This needs // to be guaranteed by the callers of UnrollRuntimeLoopRemainder. - BranchInst *LatchBR = cast(Latch->getTerminator()); - assert((LatchBR->getSuccessor(0) == LatchExit || - LatchBR->getSuccessor(1) == LatchExit) && - "one of the loop latch successors should be " - "the exit block!"); - (void)LatchBR; + assert(!L->contains(LatchBR->getSuccessor(ExitIndex)) && + "one of the loop latch successors should be the exit block!"); + // Support runtime unrolling for multiple exit blocks and multiple exiting + // blocks. + if (!LatchExit) { + assert(UseEpilogRemainder && "Multi exit unrolling is currently supported " + "unrolling with epilog remainder only!"); + LatchExit = LatchBR->getSuccessor(ExitIndex); + // TODO: Support multiple exiting blocks jumping to the `LatchExit`. This + // will need updating the logic in connectEpilog. + if (!LatchExit->getSinglePredecessor()) + return false; + SmallVector Exits; + L->getUniqueExitBlocks(Exits); + for (auto *BB : Exits) + if (BB != LatchExit) + OtherExits.push_back(BB); + } + + assert(LatchExit && "Latch Exit should exist!"); + // Use Scalar Evolution to compute the trip count. This allows more loops to // be unrolled than relying on induction var simplification. if (!SE) @@ -512,7 +537,6 @@ if (isa(TripCountSC)) return false; - BasicBlock *Header = L->getHeader(); BasicBlock *PreHeader = L->getLoopPreheader(); BranchInst *PreHeaderBR = cast(PreHeader->getTerminator()); const DataLayout &DL = Header->getModule()->getDataLayout(); @@ -654,8 +678,9 @@ // iterations. This function adds the appropriate CFG connections. BasicBlock *InsertBot = UseEpilogRemainder ? LatchExit : PrologExit; BasicBlock *InsertTop = UseEpilogRemainder ? EpilogPreHeader : PrologPreHeader; - CloneLoopBlocks(L, ModVal, CreateRemainderLoop, UseEpilogRemainder, InsertTop, - InsertBot, NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI); + Loop *remainderLoop = CloneLoopBlocks( + L, ModVal, CreateRemainderLoop, UseEpilogRemainder, InsertTop, InsertBot, + NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI); // Insert the cloned blocks into the function. F->getBasicBlockList().splice(InsertBot->getIterator(), @@ -725,6 +750,60 @@ if (Loop *ParentLoop = L->getParentLoop()) SE->forgetLoop(ParentLoop); + // Now the loop blocks are cloned and the edges between the latchExit and the + // prolog/epilog blocks are correctly setup. Also, other exiting blocks from + // the epilog 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. Then canonicalize the original and remainder loops + // to LoopSimplifyForm (i.e. having dedicated exits). + for (auto *BB : OtherExits) + for (auto &II : *BB) { + if (!isa(II)) + break; + PHINode *OldPHI = cast(&II); + SmallVector OrigLoopBlocks; + SmallVector OrigLoopValues; + unsigned oldNumOperands = OldPHI->getNumIncomingValues(); + for (unsigned i =0; i < oldNumOperands; i++){ + Value *newVal = VMap[OldPHI->getIncomingValue(i)]; + if (!newVal) { + assert(isa(OldPHI->getIncomingValue(i)) && + "VMap should exist for all values except constants!"); + newVal = OldPHI->getIncomingValue(i); + } + OldPHI->addIncoming(newVal, + cast(VMap[OldPHI->getIncomingBlock(i)])); + } + } + + // Update to LoopSimplifyForm for both original and remainder loops. We cannot + // rely on the LoopUnroll to do this because it only does canonicalization for + // parent/subloops and not the sibling loops. + if (OtherExits.size() > 0) { + // Generate dedicated exit blocks for the original loop, to preserve + // LoopSimplifyForm. + formDedicatedExitBlocks(L, DT, LI, PreserveLCSSA); + // Generate dedicated exit blocks for the remainder loop if one exists, to + // preserve LoopSimplifyForm. + if (remainderLoop) + formDedicatedExitBlocks(remainderLoop, DT, LI, PreserveLCSSA); + // NOTE! Update dominator info for the newly generated exit block(s) that contain + // edges from the dedicated exits of both the original and remainder loops. + // After the call to formDedicatedExitBlocks, the newly generated exit + // block(s) have immediate dominator set to the header of the loop. This is + // incorrect because the original and remainder loops are siblings! So, we + // need extra dominator info updates for adjacent loops that + // shared exit blocks. + SmallVector newExitsForOriginalLoop; + L->getUniqueExitBlocks(newExitsForOriginalLoop); + if (DT) + for (auto *BB : newExitsForOriginalLoop) { + // The new *dedicated* exits for the original loop have + // exactly one successor. + DT->changeImmediateDominator(BB->getSingleSuccessor(), PreHeader); + } + } + NumRuntimeUnrolled++; return true; } Index: test/Transforms/LoopUnroll/runtime-loop-multiple-exits.ll =================================================================== --- /dev/null +++ test/Transforms/LoopUnroll/runtime-loop-multiple-exits.ll @@ -0,0 +1,176 @@ +; RUN: opt < %s -loop-unroll -unroll-runtime=true -unroll-runtime-epilog=true -unroll-runtime-multi-exit=true -verify-dom-info -verify-loop-info -instcombine -S | FileCheck %s + +; test with three exiting and three exit blocks. +define void @test1(i64 %trip) { +; CHECK-LABEL: test1 +; CHECK: %niter.nsub.7 = add i64 %niter, -8 +entry: + br label %loop_header +loop_header: + %iv = phi i64 [ 0, %entry ], [ %iv_next, %loop_latch ] + br i1 undef, label %loop_latch, label %loop_exiting_bb1 +loop_exiting_bb1: + br i1 false, label %loop_exiting_bb2, label %exit1 +loop_exiting_bb2: + br i1 false, label %loop_latch, label %exit3 +exit3: + ret void +loop_latch: + %iv_next = add i64 %iv, 1 + %cmp = icmp ne i64 %iv_next, %trip + br i1 %cmp, label %loop_header, label %exit2.loopexit +exit1: + ret void +exit2.loopexit: + ret void +} + + +; test with three exiting and two exit blocks +; There are 2 values for the return that are calculated at every iteration. +; %sum.02 and %add. Both of these are incoming values for phi from every exiting +; unrolled block. +define i32 @test2(i32* nocapture %a, i64 %n) nounwind uwtable readonly { +; CHECK-LABEL: test2 +; CHECK-LABEL: for.exit2.loopexit: +; CHECK-NEXT: %retval.ph = phi i32 [ 42, %for.exiting_block ], [ %sum.02, %header ], [ %add, %for.body ], [ 42, %for.exiting_block.1 ], [ %add.1, %for.body.1 ], [ 42, %for.exiting_block.2 ], [ %add.2, %for.body.2 ], [ 42, %for.exiting_block.3 ], +; CHECK-NEXT: br label %for.exit2 +; CHECK-LABEL: for.exit2.loopexit2: +; CHECK-NEXT: %retval.ph3 = phi i32 [ 42, %for.exiting_block.epil ], [ %sum.02.epil, %header.epil ] +; CHECK-NEXT: br label %for.exit2 +; CHECK-LABEL: for.exit2: +; CHECK-NEXT: %retval = phi i32 [ %retval.ph, %for.exit2.loopexit ], [ %retval.ph3, %for.exit2.loopexit2 ] +; CHECK-NEXT: ret i32 %retval +; CHECK: %niter.nsub.7 = add i64 %niter, -8 +entry: + br label %header + +header: + %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ] + %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ] + br i1 false, label %for.exit2, label %for.exiting_block + +for.exiting_block: + %cmp = icmp eq i64 %n, 42 + br i1 %cmp, label %for.exit2, label %for.body + +for.body: ; preds = %for.body, %entry + %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv + %0 = load i32, i32* %arrayidx, align 4 + %add = add nsw i32 %0, %sum.02 + %indvars.iv.next = add i64 %indvars.iv, 1 + %exitcond = icmp eq i64 %indvars.iv.next, %n + br i1 %exitcond, label %for.end, label %header + +for.end: ; preds = %for.body, %entry + %sum.0.lcssa = phi i32 [ %add, %for.body ] + ret i32 %sum.0.lcssa + +for.exit2: + %retval = phi i32 [ %sum.02, %header ], [ 42, %for.exiting_block ] + ret i32 %retval +} + +; FIXME: Support multiple exiting blocks to the same latch exit block. +define i32 @test3(i32* nocapture %a, i64 %n) nounwind uwtable readonly { +; CHECK-LABEL: test3 +; CHECK-NOT: .unr +entry: + br label %header + +header: + %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ] + %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ] + br i1 undef, label %for.end, label %for.exiting_block + ;br i1 undef, label %for.exit2, label %for.exiting_block + +for.exiting_block: + %cmp = icmp eq i64 %n, 42 + br i1 %cmp, label %for.exit2, label %for.body + +for.body: ; preds = %for.body, %entry + %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv + %0 = load i32, i32* %arrayidx, align 4 + %add = add nsw i32 %0, %sum.02 + %indvars.iv.next = add i64 %indvars.iv, 1 + %exitcond = icmp eq i64 %indvars.iv.next, %n + br i1 %exitcond, label %for.end, label %header + +for.end: ; preds = %for.body, %entry + %sum.0.lcssa = phi i32 [ 0, %header ], [ %add, %for.body ] + ret i32 %sum.0.lcssa + +for.exit2: + ret i32 42 +} + + +; test with two exiting and three exit blocks. +; the non-latch exiting block has a switch. +define void @test4(i64 %trip, i64 %add) { +; CHECK-LABEL: test4 +; CHECK: %niter.nsub.7 = add i64 %niter, -8 +entry: + br label %loop_header + +loop_header: + %iv = phi i64 [ 0, %entry ], [ %iv_next, %loop_latch ] + %sum = phi i64 [ 0, %entry ], [ %sum.next, %loop_latch ] + br i1 undef, label %loop_latch, label %loop_exiting_bb1 + +loop_exiting_bb1: + switch i64 %sum, label %loop_latch [ + i64 24, label %exit1 + i64 42, label %exit3 + ] + +exit3: + ret void + +loop_latch: + %iv_next = add nuw nsw i64 %iv, 1 + %sum.next = add i64 %sum, %add + %cmp = icmp ne i64 %iv_next, %trip + br i1 %cmp, label %loop_header, label %exit2.loopexit + +exit1: + ret void + +exit2.loopexit: + ret void +} + +define i32 @test5(i32* nocapture %a, i64 %n) nounwind uwtable readonly { +; CHECK-LABEL: test5 +; CHECK: %niter.nsub.7 = add i64 %niter, -8 +entry: + br label %header + +header: + %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ] + %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ] + br i1 false, label %for.exit2, label %for.exiting_block + +for.exiting_block: + %cmp = icmp eq i64 %n, 42 + br i1 %cmp, label %for.exit2, label %for.body + +for.body: ; preds = %for.body, %entry + %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv + %0 = load i32, i32* %arrayidx, align 4 + %add = add nsw i32 %0, %sum.02 + %indvars.iv.next = add i64 %indvars.iv, 1 + %exitcond = icmp eq i64 %indvars.iv.next, %n + br i1 %exitcond, label %for.end, label %header + +for.end: ; preds = %for.body, %entry + %sum.0.lcssa = phi i32 [ %add, %for.body ] + ret i32 %sum.0.lcssa + +for.exit2: + %retval = phi i32 [ %sum.02, %header ], [ 42, %for.exiting_block ] + br label %exit_succ + +exit_succ: + ret i32 %retval +}