[LV] Do not fold the tail when an IV is used outside of the loop.

[rickyz]

Summary

Tail-folding does not support cases where an IV is used outside of the
loop. Before this change, the tail-folding legality check did not check
for this. In PR52335, this led to generating code with an incorrrect
index computation. Fix this by adding the missing check.

This change also adds assertions that there are no externally-used IVs
when performing tail folding.

Fixes PR52335.

[rickyz]

Here's some additional context on this bug:
https://github.com/llvm/llvm-project/issues/51677#issuecomment-991585227

I'm pretty new to this code, so please let me know anything if I've said anything inaccurate!

[fhahn]

Thanks for the patch!

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

576

While this indeed fixes the issue for the test case, I think this is not really the right fix.

IIUC AllowedExit is supposed to contain the set of instructions that are *allowed* to have outside users. So we need to check *all* instructions that have outsides users are also in this set.

But that's not what is happening at the moment. We won't check reduction or induction phis against this set in general.

For tail folding, we need to make sure that there are no users outside the loop for any induction, independent of whether they are allowed to have outside users or not (i.e. no IV is allowed to have outside users at the moment for tail folding).

A direct way to fix this would be to use the following

--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -1292,6 +1292,17 @@ bool LoopVectorizationLegality::prepareToFoldTailByMasking() {
   for (auto &Reduction : getReductionVars())
     ReductionLiveOuts.insert(Reduction.second.getLoopExitInstr());

+  for (auto &Induction : getInductionVars())
+    if (any_of(Induction.first->users(), [this](User *U) {
+      return !TheLoop->contains(cast<Instruction>(U));
+    })) {
+      LLVM_DEBUG(
+          dbgs()
+          << "LV: Cannot fold tail by masking, loop has an outside user for "
+          << *Induction.first << "\n");
+      return false;
+    }
+

Note that it might be possible to compute the final induction value even when tail-folding similar to how we compute the final reduction value, but that's only a potential follow up.

Independent of that, we may be able to allow users of the phi itself even if there is a SCEV predicate in general.

Note that the current logic does not prevent vectorization if the induction phi is used outside in the predicated case. Without tail folding, we never check if an induction phi has users outside the loop. A possible justification is the following: the phi value is guaranteed to be valid for each iteration of the loop as per the SCEV predicate, so even if it is used outside the loop, the value is only used if the predicate holds. For the incremented value this is not true. The value computed in the last iteration may be poison and should not be used outside the scope of the predicate.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

568

All changes in this file seem to be a cleanup independent of fixing the issue. If that's the case, could you submit them as a separate patch for review?

llvm/test/Transforms/LoopVectorize/pr52335.ll

6

The C example here is mostly irrelevant as the IR generated for it may change over time and depends on the flags provided and so on. The IR contains all the relevant information and is compact. Ideally the comment for the function would describe the issue. The C source is available at the bug report, which should be sufficient.

20

this should not be needed with -force-vector-width=4. To be safe, also pass -force-vector-interleave=1.

39

nit: moving the exit blocks after the loop slightly increases readability. Also the names of the blocks could be a bit more compact.

[fhahn]

Reverse ping. @rickyz do you plan on picking this up again?

[rickyz]

Apologies for the long delay, I've been a little sidetracked from LLVM contributions recently - I plan to get back to this within the next week or so - thank you for your detailed comments!

[rickyz]

Split out cleanup, address comments.

[rickyz]

Apologies for the super long delay on responding to this!

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

576

Thank you for the detailed explanation and suggestion - that makes sense to me. I'll split out the AllowedExit change into a separate change (it should become a NFC change after this fix).

[rickyz]

Rebase.

[fhahn]

LGTM, thanks!

llvm/test/Transforms/LoopVectorize/pr52335.ll

11

I think it would be good to check the full vector body to make the test more robust. It would also be good to submit the test separately before and then only have the test-diff for the change here.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3328

this is independent of the patch and could be moved out of the patch, right? If so, it would be good to split of.

[fhahn]

@rickyz I think all this needs is to pre-commit the test and then land the change. Please let me know if you have commit access. If not, I can land the patches on your behalf. In that case, please let me know which name & email I should use to set the git authorship.

[fhahn]

@rickyz I am planning on going ahead and landing this soon using the info from Phabricator.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

if (!PrimaryInduction || PhiTy == WidestIndTy)
       PrimaryInduction = Phi;
   }
 
   // Both the PHI node itself, and the "post-increment" value feeding
   // back into the PHI node may have external users.
   // We can allow those uses, except if the SCEVs we have for them rely
   // on predicates that only hold within the loop, since allowing the exit
   // currently means re-using this SCEV outside the loop (see PR33706 for more
   // details).
   if (PSE.getPredicate().isAlwaysTrue()) {
     AllowedExit.insert(Phi);
     AllowedExit.insert(Phi->getIncomingValueForBlock(TheLoop->getLoopLatch()));
   }
 
   LLVM_DEBUG(dbgs() << "LV: Found an induction variable.\n");
 }
 
bool LoopVectorizationLegality::canVectorize(bool UseVPlanNativePath) {
   // Okay! We've done all the tests. If any have failed, return false. Otherwise
   // we can vectorize, and at this point we don't have any other mem analysis
   // which may limit our maximum vectorization factor, so just return true with
   // no restrictions.
   return Result;
 }
 
 bool LoopVectorizationLegality::prepareToFoldTailByMasking() {
-
   LLVM_DEBUG(dbgs() << "LV: checking if tail can be folded by masking.\n");
 
-  SmallPtrSet<const Value *, 8> ReductionLiveOuts;
+  for (auto &Induction : getInductionVars())
+    if (any_of(Induction.first->users(), [&](User *U) {
+          return !TheLoop->contains(cast<Instruction>(U));
+        })) {
+      LLVM_DEBUG(
+          dbgs()
+          << "LV: Cannot fold tail by masking, loop has an outside user for "
+          << *Induction.first << "\n");
+      return false;
+    }
 
+  SmallPtrSet<const Value *, 8> ReductionLiveOuts;
   for (auto &Reduction : getReductionVars())
     ReductionLiveOuts.insert(Reduction.second.getLoopExitInstr());
 
   // TODO: handle non-reduction outside users when tail is folded by masking.
   for (auto *AE : AllowedExit) {
     // Check that all users of allowed exit values are inside the loop or
     // are the live-out of a reduction.
     if (ReductionLiveOuts.count(AE))

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

protected:
 
   /// A type for scalarized values in the new loop. Each value from the
   /// original loop, when scalarized, is represented by UF x VF scalar values
   /// in the new unrolled loop, where UF is the unroll factor and VF is the
   /// vectorization factor.
   using ScalarParts = SmallVector<SmallVector<Value *, 4>, 2>;
 
   /// Set up the values of the IVs correctly when exiting the vector loop.
   void fixupIVUsers(PHINode *OrigPhi, const InductionDescriptor &II,
                     Value *VectorTripCount, Value *EndValue,
                     BasicBlock *MiddleBlock, BasicBlock *VectorHeader);
 
   /// Handle all cross-iteration phis in the header.
   void fixCrossIterationPHIs(VPTransformState &State);
 
   /// Create the exit value of first order recurrences in the middle block and
   /// update their users.
void InnerLoopVectorizer::fixupIVUsers(PHINode *OrigPhi,
 
   DenseMap<Value *, Value *> MissingVals;
 
   // An external user of the last iteration's value should see the value that
   // the remainder loop uses to initialize its own IV.
   Value *PostInc = OrigPhi->getIncomingValueForBlock(OrigLoop->getLoopLatch());
   for (User *U : PostInc->users()) {
     Instruction *UI = cast<Instruction>(U);
-    if (!OrigLoop->contains(UI)) {
-      assert(isa<PHINode>(UI) && "Expected LCSSA form");
-      MissingVals[UI] = EndValue;
-    }
-  }
+    if (OrigLoop->contains(UI))
+      continue;
+
+    assert(!Cost->foldTailByMasking() && "Unexpected external use of IV");
 
   // An external user of the penultimate value need to see EndValue - Step.
   // The simplest way to get this is to recompute it from the constituent SCEVs,
   // that is Start + (Step * (CRD - 1)).
   for (User *U : OrigPhi->users()) {
     auto *UI = cast<Instruction>(U);
     if (!OrigLoop->contains(UI)) {
       assert(isa<PHINode>(UI) && "Expected LCSSA form");
+      assert(!Cost->foldTailByMasking() && "Unexpected external use of IV");
 
       IRBuilder<> B(MiddleBlock->getTerminator());
 
       // Fast-math-flags propagate from the original induction instruction.
       if (II.getInductionBinOp() && isa<FPMathOperator>(II.getInductionBinOp()))
         B.setFastMathFlags(II.getInductionBinOp()->getFastMathFlags());
 
       Value *CountMinusOne = B.CreateSub(

llvm/test/Transforms/LoopVectorize/pr52335.ll

+; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S | FileCheck %s
+
+define dso_local i32 @test(i32* nocapture %arr, i64 %n) {
+; The vectorizer should refuse to fold the tail by masking because
+; %conv is used outside of the loop. Test for this by checking that
+; %n.vec, the vector trip count, is rounded down to the next multiple of
+; 4. If folding the tail, it would have been rounded up instead.
+;
+; CHECK-LABEL: vector.ph:
+; CHECK-NEXT:    %n.mod.vf = urem i64 [[TC:%.*]], 4
+; CHECK-NEXT:    %n.vec = sub i64 [[TC]], %n.mod.vf
+entry:
+  %cmp1 = icmp ugt i64 %n, 1
+  br i1 %cmp1, label %preheader, label %done
+
+preheader:
+  br label %loop
+
+loop:
+  %conv = phi i64 [ %conv2, %loop ], [ 1, %preheader ]
+  %i = phi i8 [ %inc, %loop ], [ 1, %preheader ]
+  %sub = add nsw i64 %conv, -1
+  %ptr = getelementptr inbounds i32, i32* %arr, i64 %sub
+  store i32 65, i32* %ptr, align 4
+  %inc = add i8 %i, 1
+  %conv2 = zext i8 %inc to i64
+  %cmp2 = icmp ult i64 %conv2, %n
+  br i1 %cmp2, label %loop, label %load_val, !llvm.loop !0
+
+load_val:
+  %final = phi i64 [ %conv, %loop ]
+  %ptr2 = getelementptr inbounds i32, i32* %arr, i64 %final
+  %val = load i32, i32* %ptr2, align 4
+  br label %done
+
+done:
+  %value = phi i32 [ %val, %load_val ], [ 0, %entry ]
+  ret i32 %value
+
+}
+
+!0 = distinct !{!0, !1, !2, !3}
+!1 = !{!"llvm.loop.unroll.disable"}
+!2 = !{!"llvm.loop.vectorize.predicate.enable", i1 true}
+!3 = !{!"llvm.loop.vectorize.enable", i1 true}