[IndVarSimplify] Rewrite loop exit values with their initial values from loop preheader

[• chenli]

Summary

This is a revised version of D13974, and the following quoted summary are from D13974

"This patch adds support to check if a 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. If there is an induction variable used in that exit path whose value has not been updated, it will keep its initial value passing from loop preheader. We can therefore rewrite the exit value with
its initial value. This will help remove phis created by LCSSA and enable other optimizations like loop unswitch."

D13974 was committed but failed one lnt test. The bug was that we only checked the condition from loop exit's incoming block was a loop invariant. But there could be another condition from loop header to that incoming block not being a loop invariant. This would produce miscompiled code.

This patch fixes the issue by checking if the incoming block is loop header, and if not, don't perform the rewrite. The could be further improved by recursively checking all conditions leading to loop exit block, but I'd like to check in this simple version first and improve it with future patches.

[• chenli]

The original diff file was from D13974 for the purpose of comparison. This is the updated diff file.

lib/Transforms/Scalar/IndVarSimplify.cpp

728

This is the fix to the problem mentioned in summary.

[sanjoy]

lgtm with minor nits inline

lib/Transforms/Scalar/IndVarSimplify.cpp

719

Nit: LLVM style is BasicBlock::iterator BBI = ExitBB->begin(); or BasicBlock::iterator Begin = ExitBB->begin();

745

I don't think you have to special case this -- isLoopInvariant takes a Value * and returns true for non instructions.

764

Since you're calling getBasicBlockIndex, why not save the index and re-use it here?

lib/Transforms/Scalar/IndVarSimplify.cpp

private:
 
   void handleFloatingPointIV(Loop *L, PHINode *PH);
   void rewriteNonIntegerIVs(Loop *L);
 
   void simplifyAndExtend(Loop *L, SCEVExpander &Rewriter, LoopInfo *LI);
 
   bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet);
   void rewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter);
+  void rewriteFirstIterationLoopExitValues(Loop *L);
 
   Value *linearFunctionTestReplace(Loop *L, const SCEV *BackedgeTakenCount,
                                    PHINode *IndVar, SCEVExpander &Rewriter);
 
   void sinkUnusedInvariants(Loop *L);
 
   Value *expandSCEVIfNeeded(SCEVExpander &Rewriter, const SCEV *S, Loop *L,
                             Instruction *InsertPt, Type *Ty);
for (const RewritePhi &Phi : RewritePhiSet) {
     }
   }
 
   // The insertion point instruction may have been deleted; clear it out
   // so that the rewriter doesn't trip over it later.
   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<BasicBlock *, 8> 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<PHINode>(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<BranchInst>(TermInst)) {
+          // Must be a conditional branch, otherwise the block
+          // should not be in the loop.
+          Cond = BI->getCondition();
+        } else if (auto *SI = dyn_cast<SwitchInst>(TermInst))
+          Cond = SI->getCondition();
+        else
+          continue;
+
+        // All non-instructions are loop-invariant.
+        if (isa<Instruction>(Cond) && !L->isLoopInvariant(Cond))
+            continue;
+
+        auto *ExitVal =
+            dyn_cast<PHINode>(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.
 bool IndVarSimplify::canLoopBeDeleted(
     Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet) {
 
   BasicBlock *Preheader = L->getLoopPreheader();
   // If there is no preheader, the loop will not be deleted.
if (Instruction *Inst =
       RecursivelyDeleteTriviallyDeadInstructions(Inst, TLI);
 
   // The Rewriter may not be used from this point on.
 
   // Loop-invariant instructions in the preheader that aren't used in the
   // 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);
 
   // Check a post-condition.
   assert(L->isRecursivelyLCSSAForm(*DT) && "Indvars did not preserve LCSSA!");
 
   // Verify that LFTR, and any other change have not interfered with SCEV's
   // ability to compute trip count.

test/Transforms/IndVarSimplify/rewrite-loop-exit-value.ll

+; 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
+}
+