Index: lib/Transforms/Scalar/IndVarSimplify.cpp =================================================================== --- lib/Transforms/Scalar/IndVarSimplify.cpp +++ lib/Transforms/Scalar/IndVarSimplify.cpp @@ -133,6 +133,7 @@ bool canLoopBeDeleted(Loop *L, SmallVector &RewritePhiSet); void rewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter); + void rewriteFirstIterationLoopExitValues(Loop *L); Value *linearFunctionTestReplace(Loop *L, const SCEV *BackedgeTakenCount, PHINode *IndVar, SCEVExpander &Rewriter); @@ -695,6 +696,78 @@ Rewriter.clearInsertPoint(); } +//===---------------------------------------------------------------------===// +// rewriteFirstIterationLoopExitValues: Rewrite loop exit values if we know +// they will exit at the first iteration. +//===---------------------------------------------------------------------===// + +/// Check to see if this loop has loop invariant conditions which lead to loop +/// exits. If so, we know that if the exit path is taken, it is at the first +/// loop iteration. This lets us predict exit values of PHI nodes that live in +/// loop header. +void IndVarSimplify::rewriteFirstIterationLoopExitValues(Loop *L) { + // Verify the input to the pass is already in LCSSA form. + assert(L->isLCSSAForm(*DT)); + + SmallVector ExitBlocks; + L->getUniqueExitBlocks(ExitBlocks); + + for (auto *ExitBB : ExitBlocks) { + BasicBlock::iterator begin = ExitBB->begin(); + // If there are no more PHI nodes in this exit block, then no more + // values defined inside the loop are used on this path. + while (auto *PN = dyn_cast(begin++)) { + for (unsigned IncomingValIdx = 0, e = PN->getNumIncomingValues(); + IncomingValIdx != e; ++IncomingValIdx) { + auto *IncomingBB = PN->getIncomingBlock(IncomingValIdx); + + // We currently only support loop exits from loop header. If the + // incoming block is not loop header, we need to recursively check + // all conditions starting from loop header are loop invariants. + // Additional support might be added in the future. + if (IncomingBB != L->getHeader()) + continue; + + // Get condition that leads to the exit path. + auto *TermInst = IncomingBB->getTerminator(); + + Value *Cond = nullptr; + if (auto *BI = dyn_cast(TermInst)) { + // Must be a conditional branch, otherwise the block + // should not be in the loop. + Cond = BI->getCondition(); + } else if (auto *SI = dyn_cast(TermInst)) + Cond = SI->getCondition(); + else + continue; + + // All non-instructions are loop-invariant. + if (isa(Cond) && !L->isLoopInvariant(Cond)) + continue; + + auto *ExitVal = + dyn_cast(PN->getIncomingValue(IncomingValIdx)); + + // Only deal with PHIs. + if (!ExitVal) + continue; + + // If ExitVal is a PHI on the loop header, then we know its + // value along this exit because the exit can only be taken + // on the first iteration. + auto *LoopPreheader = L->getLoopPreheader(); + assert(LoopPreheader && "Invalid loop"); + if (ExitVal->getBasicBlockIndex(LoopPreheader) != -1) { + assert(ExitVal->getParent() == L->getHeader() && + "ExitVal must be in loop header"); + PN->setIncomingValue(IncomingValIdx, + ExitVal->getIncomingValueForBlock(LoopPreheader)); + } + } + } + } +} + /// Check whether it is possible to delete the loop after rewriting exit /// value. If it is possible, ignore ReplaceExitValue and do rewriting /// aggressively. @@ -2154,6 +2227,11 @@ // loop may be sunk below the loop to reduce register pressure. sinkUnusedInvariants(L); + // rewriteFirstIterationLoopExitValues does not rely on the computation of + // trip count and therefore can further simplify exit values in addition to + // rewriteLoopExitValues. + rewriteFirstIterationLoopExitValues(L); + // Clean up dead instructions. Changed |= DeleteDeadPHIs(L->getHeader(), TLI); Index: test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll =================================================================== --- /dev/null +++ test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll @@ -0,0 +1,63 @@ +; RUN: opt -indvars -instcombine -S < %s | FileCheck %s + +;; Test that loop's exit value is rewritten to its initial +;; value from loop preheader +define i32 @test1(i32* %var) { +; CHECK-LABEL: @test1 +entry: + %cond = icmp eq i32* %var, null + br label %header + +header: + %phi_indvar = phi i32 [0, %entry], [%indvar, %loop] + br i1 %cond, label %loop, label %exit + +loop: + %indvar = add i32 %phi_indvar, 1 + br label %header + +exit: +; CHECK: ret i32 0 + ret i32 %phi_indvar +} + +;; Test that we can not rewrite loop exit value if it's not +;; a phi node (%indvar is an add instruction in this test). +define i32 @test2(i32* %var) { +; CHECK-LABEL: @test2 +entry: + %cond = icmp eq i32* %var, null + br label %header + +header: + %phi_indvar = phi i32 [0, %entry], [%indvar, %header] + %indvar = add i32 %phi_indvar, 1 + br i1 %cond, label %header, label %exit + +exit: +; CHECK: ret i32 %indvar + ret i32 %indvar +} + +;; Test that we can not rewrite loop exit value if the condition +;; is not in loop header. +define i32 @test3(i32* %var) { +; CHECK-LABEL: @test3 +entry: + %cond1 = icmp eq i32* %var, null + br label %header + +header: + %phi_indvar = phi i32 [0, %entry], [%indvar, %header], [%indvar, %body] + %indvar = add i32 %phi_indvar, 1 + %cond2 = icmp eq i32 %indvar, 10 + br i1 %cond2, label %header, label %body + +body: + br i1 %cond1, label %header, label %exit + +exit: +; CHECK: ret i32 %phi_indvar + ret i32 %phi_indvar +} +