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

[• chenli]

Summary

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.

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

137

I'm not sure that RewriteLoopExitValuesWithoutBECount is the best name for this function. The transform is only incidentally being done over a loop for which we have no backedge count.

How about RewriteFirstIterationLoopExitValues, since we're really using the fact that the exit, if taken, will only be taken in the first ever loop iteration?

715

I think "updated" is misleading when talking about SSA. Why not keep the comment brief and simple, something like:

/// Check to see if this loop has exits conditional on loop invariant values.  If so, we know that if the exit path is taken, it is taken on the first iteration.  This lets us predict the value of well behaved induction variables.

718

I'm not sure if "As a side effect, the phi node generated by LCSSA is
/// no longer needed." is accurate -- can't LCSSA phis have multiple incoming values?

[edit] See longer discussion below.

727

Nit: spacing.

I'll suggest just running the diff through clang-format before submitting.

734

Use auto * here.

735

I think this is misleading -- can't LCSSA phis have multiple incoming values? I'm thinking of cases like these:

for (...) {
  t = ...;
  if (..) goto x;
  ..
  if (..) goto x;
}
...
x:
  print(t)

t dominates x so in "normal" code you won't need a PHI, but if you do want to insert a PHI for t then you'll need a phi [x, x]. Actually I'm not even sure if an "LCSSA PHI" is a well defined concept.

In any case, I don't think there is any need to special case LCSSA PHI's here. You could special case single entry PHIs if you really wanted to, but I doubt that restriction will make your change any simpler.

742

This is mostly stylistic, but prefer using auto * on the LHS of a dyn_cast; since the type is obvious.

752

What if Cond is a LoadInst loading from an i1 *?

758

How do you know that the 0th incoming value to PN is an induction PHI node?

770

I don't quite follow the logic of when InductionPHI->getBasicBlockIndex(LoopPreheader) can be -1. The preheader is the only block outside the loop branching into the header, and an induction variable (which will be a PHI node on the loop header) should have to have one incoming value for the "first" iteration which will have to be the value incoming from the preheader.

If you see failures without this check, I suspect this is because, as I mentioned in the previous comment, PN->getIncomingValue(0) need not be an induction PHI node.

[• chenli]

lib/Transforms/Scalar/IndVarSimplify.cpp

735

You are correct. For some reason I thought LCSSA PHIs were always single entry PHIs so I used the check to filter non LCSSA PHIs. I will fix it.

752

I think I'll just use Loop::isLoopInvariant(const Value *V) instead.

758

This is basically because of the previous assumption that we only have single entry PHI here. So if there is an induction PHI node, it should be 0th. The later

InductionPHI->getBasicBlockIndex(LoopPreheader) != -1

is used to filter non-induction PHI node.
The name InductionPHI is confusing because it's not necessary an indication PHI if it fails the check.
I will refactor the code to make it cleaner.

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

758

In any case, we don't really care about the PHI being an induction variable (i.e. a integer that increases / decreases by some amount every iteration) -- whatever the value incoming from the preheader -> header edge is will be the exit value of the PHI. So perhaps it makes sense to not call the value InductionPHI?

[• chenli]

lib/Transforms/Scalar/IndVarSimplify.cpp

758

Yes, I will definitely pick a better name :)

[• chenli]

Update patch w.r.t Sanjoy's request.

[sanjoy]

Mostly minor stylistic comments inline.

lib/Transforms/Scalar/IndVarSimplify.cpp

715

Now that we're not longer special casing induction variable, please change this comment to be something like "lets us predict values of PHI nodes that live on the loop header".

720

I think this is a false promise :) There are restrictions around other
cases, e.g. (using a terrible LLVM and C mashup):

 for (init; cond; incr) {
  ..
  if (c) {
   ..
   br merge;
  } else {
   ..
   br merge;
  }
 merge:
  %val = phi [ 0, %if ], [ 1, %else ]
  if (loop_invariant_cond)
    br exit
  ...
}

exit:
  %lcssa = phi [ %val ]

You cannot predict the value of %val here.

I'd just remove this NOTE for now.

733

I'd suggest doing one of the following

• rename i to something more descriptive (like incomingValIdx)

OR

• extract out the loop body to a lambda called ProcessIncomingValue or something like that.

751

Super minor: do we actually need a line break here?

754

I'd remove the Currently -- this optimization is fundamentally intended to work with PHI nodes.

760

I think (if the exit value is in a joint block of the oop then it will not be qualified) bit is somewhat redundant.

I'd just say "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" and leave it at that. No point in discussing the invalid cases.

763

Add an assert here ExitVal lives on the loop header.

2132

This may be a good place to state that RewriteFirstIterationLoopExitValues does not depend on us being able to compute the trip count of L, and therefore can get cases that RewriteLoopExitValues does not. You can also put this "factoid" in the declaration of RewriteFirstIterationLoopExitValues

[sanjoy]

(Nothing new, just to move the review out of my "blocked on you" queue).

[• chenli]

Update patch w.r.t Sanjoy's comments.

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

729

I know I suggested this, but LLVM style is IncomingValIdx.

732

Do you not need to check that PN->getIncomingBlock(incomingValIdx) is inside the loop?

736

Note: this transform is also valid for unconditional branches (since br label %foo is equivalent to br true label %foo label %foo).

Actually, thinking about it, I think we can pretty much assert(BI->isConditional()) (IOW you don't need the check since BI->getCondition() will have the same assert it in). (Please verify this) I think you cannot have an unconditional branch from a block B in the loop to a loop exit since then B will not be connected to the header (i.e. there won't be a path from B to the loop header) and hence B won't be part of the loop leading to a contradiction.

759

I'd put an assert(LoopPreaheader) here.

[• chenli]

lib/Transforms/Scalar/IndVarSimplify.cpp

736

I think you can not have an unconditional branch in B or B will be the loop exit. I will remove the check and run some test through the test cases.

[• chenli]

Update patch w.r.t Sanjoy's comments.

[sanjoy]

lgtm with comments addressed

lib/Transforms/Scalar/IndVarSimplify.cpp

743

Minor: remove the braces { form the else if.

747

]I'd change this to isLoopInvariant for now -- make LoopInvariant may hoist arithmetic etc. over control flow and this could be a regression e.g. in cases like

for ( ... )
  if (rare_condition) {
    int x = loop_invariant * 3;
    if (x == 42) then exit;
  }

hoisting loop_invariant * 3 to the loop preheader will be a net regression.

Sorry for not bringing this up earlier.

Also, note that both isLoopInvariant and makeLoopInvariant take a Value so you don't need to manually cast to an Instruction.

[• chenli]

lib/Transforms/Scalar/IndVarSimplify.cpp

747

By regression, do you mean by register pressure?

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

747

Register pressure also, but I was trying to say you'll just compute values that you don't need to compute (e.g. in the above example if rare_condition is always false at runtime then the multiplication in the loop preheader was useless).

[• chenli]

lib/Transforms/Scalar/IndVarSimplify.cpp

747

Makes sense. I will modify and submit. Thanks!

lib/Transforms/Scalar/IndVarSimplify.cpp

private:
 
   void HandleFloatingPointIV(Loop *L, PHINode *PH);
   void RewriteNonIntegerIVs(Loop *L);
 
   void SimplifyAndExtend(Loop *L, SCEVExpander &Rewriter, LPPassManager &LPM);
 
   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);
+        if (!L->contains(IncomingBB))
+          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();
+        }
+
+        // All non-instructions are loop-invariant.
+        if (auto *CondInst = dyn_cast<Instruction>(Cond))
+          if (!L->makeLoopInvariant(CondInst, Changed))
+            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->isLCSSAForm(*DT) &&
          "Indvars did not leave the loop in lcssa form!");
 
   // 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:
+ br label %header
+
+header:
+ %phi_indvar = phi i32 [0, %entry], [%indvar, %loop]
+ %cond = icmp eq i32* %var, null 
+ 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 inner loop's exit value is first rewritten to outer
+;; loop's induction variable, and then further rewritten to a 
+;; constant when process outer loop.
+define i32 @test2(i32* %var) {
+; CHECK-LABEL: @test2
+entry:
+ br label %outer_header
+
+outer_header:
+ %phi_outer = phi i32 [0, %entry], [%indvar_outer, %inner_exit]
+ br label %inner_header
+ 
+inner_header:
+ %phi_inner = phi i32 [%phi_outer, %outer_header], [%indvar_inner, %loop]
+ %cond1 = icmp eq i32* %var, null 
+ br i1 %cond1, label %loop, label %exit
+
+loop:
+ %indvar_inner = add i32 %phi_inner, 1
+ %cond2 = icmp eq i32 %indvar_inner, 10
+ br i1 %cond2, label %inner_header, label %inner_exit
+ 
+inner_exit:
+ %indvar_outer = add i32 %phi_outer, 1
+ br label %outer_header
+
+exit:
+;; %phi_inner is first rewritten to %phi_outer
+;; and then %phi_outer is rewritten to 0
+ %ret_val = add i32 %phi_inner, %phi_outer
+; CHECK: ret i32 0
+ ret i32 %ret_val
+}
+
+;; 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 @test3(i32* %var) {
+; CHECK-LABEL: @test3
+entry:
+ br label %header
+
+header:
+ %phi_indvar = phi i32 [0, %entry], [%indvar, %header]
+ %indvar = add i32 %phi_indvar, 1
+ %cond = icmp eq i32* %var, null 
+ br i1 %cond, label %header, label %exit
+
+exit:
+; CHECK: ret i32 %indvar
+ ret i32 %indvar
+}
 No newline at end of file

test/Transforms/LoopUnroll/scevunroll.ll

 ; iteration via the early exit. So loop unrolling cannot assume that
 ; the loop latch's exit count of zero is an upper bound on the number
 ; of iterations.
 ;
 ; CHECK-LABEL: @nsw_latch(
 ; CHECK: for.body:
 ; CHECK: %b.03 = phi i32 [ 0, %entry ], [ %add, %for.cond ]
 ; CHECK: return:
-; CHECK: %b.03.lcssa = phi i32 [ %b.03, %for.body ], [ %b.03, %for.cond ]
+; CHECK: %b.03.lcssa = phi i32 [ %b.03, %for.body ], [ 0, %for.cond ]
 define void @nsw_latch(i32* %a) nounwind {
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %for.cond, %entry
   %b.03 = phi i32 [ 0, %entry ], [ %add, %for.cond ]
   %tobool = icmp eq i32 %b.03, 0
   %add = add nsw i32 %b.03, 8