[LV] Avoid vectorizing first order recurrences when phi used outside loop

[anna]

Summary

In the recurrence logic, there is an assumption that we vectorize such that the
last element in the vector will be the one extracted to pass into the scalar
remainder loop. However, this is not true when there is a phi used outside the loop.
In such a case, we need the value from the second last iteration (i.e. the phi value), not the last iteration (which would be the phi update). I've added a test case for
this as well. Also see PR32396.

A follow up patch would generate the correct code gen for such cases, and turn this vectorization on.

[anna]

@mssimpso : Hi Mathew, is there a reason to have the first order recurrence test (https://reviews.llvm.org/D16197) in the AArch folder? Could we move this out, since the logic is shared by all platforms? Thanks

[mssimpso]

In D31910#723084, @anna wrote:

@mssimpso : Hi Mathew, is there a reason to have the first order recurrence test (https://reviews.llvm.org/D16197) in the AArch folder? Could we move this out, since the logic is shared by all platforms? Thanks

Yes, those tests should definitely be moved to the common directory. Thanks! Feel free to go ahead and do so if you want. I will try and take a look at your patch soon.

[mssimpso]

Hi Anna,

I think we can probably solve this without putting a limitation on the trip count. But I'm wondering if it would be better for now to just avoid vectorizing this case, at least until we can work out a more robust code generation strategy for the first-order recurrences. Why not change RecurrenceDescriptor::isFirstOrderRecurrence to ensure the phi has no users outside the loop? I think that's the underlying issue, right? We're currently assuming that only the phi update can be used externally. But in your test case, it's the phi that's used externally. Just to make sure I understand things correctly, after vectorization we end up with something like the following for the external use:

%r.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %vector.recur.extract, %middle.block ]

So if the scalar loop is executed, %r.lcssa is the phi, but if the trip count is evenly divided and the scalar loop is not executed, %r.lcssa is the phi's last update. And we currently assume we'll only encounter something like:

%r.lcssa = phi i32 [ %update, %for.body ], [ %vector.recur.extract, %middle.block ]

where it's the update that's used externally. I think the solution would eventually be something like:

middle.block:
  %vector.recur.extract = extractelement <4 x i32> %vector.update, i32 3
  %new.extract = extractelement <4 x i32> %vector.recur, i32 3
scalar.ph:
  %scalar.recur.init = phi i32 [ %vector.recur.extract, %middle.block ], ...
for.end:
  %r.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %new.extract, %middle.block ]

where we add an additional extract for the last element of the vector phi, if the phi is used externally. What do you think? We could disable the external use for now to fix the correctness issue and then re-enable it once we get the code generation right.

[anna]

Hi Mathew,

In D31910#723721, @mssimpso wrote:

Hi Anna,

I think we can probably solve this without putting a limitation on the trip count. But I'm wondering if it would be better for now to just avoid vectorizing this case, at least until we can work out a more robust code generation strategy for the first-order recurrences. Why not change RecurrenceDescriptor::isFirstOrderRecurrence to ensure the phi has no users outside the loop? I think that's the underlying issue, right? We're currently assuming that only the phi update can be used externally. But in your test case, it's the phi that's used externally. Just to make sure I understand things correctly, after vectorization we end up with something like the following for the external use:

%r.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %vector.recur.extract, %middle.block ]

So if the scalar loop is executed, %r.lcssa is the phi, but if the trip count is evenly divided and the scalar loop is not executed, %r.lcssa is the phi's last update. And we currently assume we'll only encounter something like:

%r.lcssa = phi i32 [ %update, %for.body ], [ %vector.recur.extract, %middle.block ]

where it's the update that's used externally.

Yes, that's right. The problem only occurs when we need the phi and not it's update at the end of loop. But there's 2 checks if we want to vectorize all cases except for the incorrect code gen cases:

1. a phi node is used externally, and
2. that phi node is not the loop termination condition.

If the phi node is the loop terminating condition, we would get the correct value for the phi even if we vectorize the loop completely (Trip count% step= 0).
This is the case (for VF = 4):

for.body:
  %inc.phi = phi i32 [ 0, %entry ], [ %inc, %for.body ]
  %val.phi = phi i32 [ 0, %entry ], [ %addx, %for.body ]
  %inc = add nsw i32 %inc.phi, 1
  %bc = zext i32 %inc.phi to i64.
  %addx = add nuw nsw i32 %inc.phi, %x
  %cmp = icmp eq i32 %inc.phi, 95. 
  br i1 %cmp, label %for.end, label %for.body

for.end:
  ret i32 %inc.phi

I think the solution would eventually be something like:

middle.block:
  %vector.recur.extract = extractelement <4 x i32> %vector.update, i32 3
  %new.extract = extractelement <4 x i32> %vector.recur, i32 3
scalar.ph:
  %scalar.recur.init = phi i32 [ %vector.recur.extract, %middle.block ], ...
for.end:
  %r.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %new.extract, %middle.block ]

where we add an additional extract for the last element of the vector phi, if the phi is used externally. What do you think? We could disable the external use for now to fix the correctness issue and then re-enable it once we get the code generation right.

Actually, this is a good idea, it would have a lesser chance of causing any perf regression than the limiting of trip count. The reason I went with trip count limitation was I was thinking of this bug as extracting from the "correct index" i.e. the last or the second last, instead of generating 2 extracts from the phi and it's last update. And currently, we dont have a way of specifying a runtime index for the extract operation.
With the code gen change there would also be a change to the cost model to consider this added shuffle instruction.

I'll work on the disabling this part as the first step (and work on the robust code gen as a follow up change). Thanks for the input!

[anna]

Updated diff to avoid vectorizing first order recurrences when phi used outside the loop
and the phi is not part of the loop terminating condition.
Based on Mathew's comments above, we have a better way to fix this rather than limiting the loop trip
count in the vectorized case.

[anna]

Change AArch64 test case that already showcases a failing test case.
We should not be vectorizing the loop in this test case.

[anna]

@mssimpso: Hi Matt, I've updated the diff to avoid vectorizing the loop when it generates incorrect code gen. Also, there's an AArch64 test case that already showcases this exact problem: generate incorrect code gen in first order recurrence. Please take a look. It's also aptly named phi_fail2!

[mssimpso]

lib/Transforms/Utils/LoopUtils.cpp

556–559 ↗

(On Diff #94877)

I don't think I understand this. Why do we need to check if the phi is used by the back-edge condition? If a phi is used by the back-edge condition, we should probably recognize it as an induction variable, not a first-order recurrence. I think you can remove this, and just check that the phi has no external users.

test/Transforms/LoopVectorize/first-order-recurrence.ll

390–401

In this test, "%inc.phi" is the induction variable and "%val.phi" is the potential first-order recurrence. We already allow induction variables to have external users, so we don't need to worry about this case. I think this test will pass without your patch.

[mssimpso]

In D31910#724159, @anna wrote:

@mssimpso: Hi Matt, I've updated the diff to avoid vectorizing the loop when it generates incorrect code gen. Also, there's an AArch64 test case that already showcases this exact problem: generate incorrect code gen in first order recurrence. Please take a look. It's also aptly named phi_fail2!

Yeah, that's true. Thanks very much for sorting this out!

[anna]

remove the check for IV having outside users. That is already handled (i.e. vectorized correctly) by the vectorizer.

[anna]

lib/Transforms/Utils/LoopUtils.cpp

556–559 ↗

(On Diff #94877)

removed this, and the test case corresponding to this passes as required.

test/Transforms/LoopVectorize/first-order-recurrence.ll

390–401

yes, thats right. Thanks for the clarification.

[mssimpso]

The added @extract_last_iteration test isn't really checking anything - it get's completely simplified by InstCombine, and the first-order recurrence isn't even used outside the loop so it's not affected by the patch. I would just remove it (you might have already meant to do this). Otherwise, this patch LGTM. Thanks!

[anna]

In D31910#724195, @mssimpso wrote:

The added @extract_last_iteration test isn't really checking anything - it get's completely simplified by InstCombine, and the first-order recurrence isn't even used outside the loop so it's not affected by the patch. I would just remove it (you might have already meant to do this). Otherwise, this patch LGTM. Thanks!

Yes, done.

[anna]

FYI: The disabling patch is checked in. I'm starting work on the robust code gen for this case.

lib/Transforms/Vectorize/LoopVectorize.cpp

protected:
 
   /// Generate a shuffle sequence that will reverse the vector Vec.
   virtual Value *reverseVector(Value *Vec);
 
   /// Returns (and creates if needed) the original loop trip count.
   Value *getOrCreateTripCount(Loop *NewLoop);
 
   /// Returns (and creates if needed) the trip count of the widened loop.
-  Value *getOrCreateVectorTripCount(Loop *NewLoop);
+  Value *getOrCreateVectorTripCount(Loop *NewLoop, Loop *OrigLoop  = nullptr);
 
   /// Emit a bypass check to see if the trip count would overflow, or we
   /// wouldn't have enough iterations to execute one vector loop.
   void emitMinimumIterationCountCheck(Loop *L, BasicBlock *Bypass);
   /// Emit a bypass check to see if the vector trip count is nonzero.
   void emitVectorLoopEnteredCheck(Loop *L, BasicBlock *Bypass);
   /// Emit a bypass check to see if all of the SCEV assumptions we've
   /// had to make are correct.
Value *InnerLoopVectorizer::getOrCreateTripCount(Loop *L) {
   if (TripCount->getType()->isPointerTy())
     TripCount =
         CastInst::CreatePointerCast(TripCount, IdxTy, "exitcount.ptrcnt.to.int",
                                     L->getLoopPreheader()->getTerminator());
 
   return TripCount;
 }
 
-Value *InnerLoopVectorizer::getOrCreateVectorTripCount(Loop *L) {
+Value *InnerLoopVectorizer::getOrCreateVectorTripCount(Loop *L,
+                                                       Loop *OrigLoop) {
   if (VectorTripCount)
     return VectorTripCount;
 
   Value *TC = getOrCreateTripCount(L);
   IRBuilder<> Builder(L->getLoopPreheader()->getTerminator());
 
   // Now we need to generate the expression for the part of the loop that the
   // vectorized body will execute. This is equal to N - (N % Step) if scalar
   // iterations are not required for correctness, or N - Step, otherwise. Step
   // is equal to the vectorization factor (number of SIMD elements) times the
   // unroll factor (number of SIMD instructions).
   Constant *Step = ConstantInt::get(TC->getType(), VF * UF);
   Value *R = Builder.CreateURem(TC, Step, "n.mod.vf");
 
+  assert(OrigLoop && "during first calculation of VectorTripCount, the orig "
+                     "loop should be passed in!");
+
+  // If the loop terminating condition uses the phi value at loop entry rather
+  // than the
+  // induction variable at the current iteration, and 
+  // extract the result at the last element of the vectorized phi (as done in
+  // fixFirstOrderRecurrence), we may get an incorrect result.
+  // Note that this
+  // problem only arises when both these conditions occur:
+  // 1. TC evenly divides Step, and the scalar loop is not run.
+  // 2. Some Phi value which is a relation based on the IV phi, is
+  // extracted out of the loop (this is the firstOrderRecurrenceRelation).
+  // For correctness, the vectorized part of the loop should be TC - Step times
+  // and the remaining iterations should be run in the scalar remainder loop.
+  // Example IR:
+  // loop:
+  // %var = phi i32 [ %init, %ph ], [ %sum, %loop ]
+  // %indvar = phi i32 [ 0, %ph ], [ %indvar.next, %loop ]
+  // %sum = add i32 %indvar, %a
+  // %indvar.next = add i32 %indvar, 1
+  // %c11 = icmp slt i32 %indvar, 95
+  // br i1 %c11, label %exit, label %loop
+  // exit:
+  // %var.lcssa = phi i32 [ %var, %loop ]
+  // ret i32 %var.lcssa
+  // Here although the trip count is 96, if we run the vectorized loop
+  // with Step = 4, %var will be extracted as 95 + a since we extract the
+  // last element of the vectorized %var, whereas it should be 94 +
+  // a.
+  // This is because the exit block uses the `var` value (calculated at previous
+  // iteration) and not the `sum` value
+  // calculated at current iteration of loop.
+  auto IsSecondLastIterationExtracted = [&]() {
+    auto *ExitBlock = OrigLoop->getExitBlock();
+    auto *HeaderBlock = OrigLoop->getHeader();
+    auto *Cmp = dyn_cast<CmpInst>(HeaderBlock->getTerminator()->getOperand(0));
+    if (!Cmp)
+      return false;
+    // Confirmed that we are effectively running TC - 1 number of iterations.
+    bool LoopTermConditionUsesPhi = any_of(
+        Cmp->operands(), [&](Use &U) { return Legal->isInductionVariable(U); });
+    if (!LoopTermConditionUsesPhi)
+      return false;
+    for (auto &I : *HeaderBlock) {
+      if (!isa<PHINode>(I))
+        break;
+      // We only care about phi nodes that depend on the IV (i.e. secondary
+      // induction variables).
+      if ((Legal->getPrimaryInduction()) == &(cast<PHINode>(I)))
+        continue;
+      for (auto *U : I.users()) {
+        if (auto *PhiUse = dyn_cast<Instruction>(U))
+          // At this stage, we know that the value for the secondary IV 
+          // should be the second last element in the vectorized version (since
+          // we are using the phi).
+          // However, we don't have the capability in vector.extractelement to
+          // choose between secondlast and last element. So, we conservatively
+          // assume that the secondlastiteration will be used and generate the
+          // loop accordingly.
+          if (PhiUse->getParent() == ExitBlock)
+            return true;
+      }
+    }
+    return false;
+  };
+
   // If there is a non-reversed interleaved group that may speculatively access
   // memory out-of-bounds, we need to ensure that there will be at least one
   // iteration of the scalar epilogue loop. Thus, if the step evenly divides
   // the trip count, we set the remainder to be equal to the step. If the step
   // does not evenly divide the trip count, no adjustment is necessary since
   // there will already be scalar iterations. Note that the minimum iterations
   // check ensures that N >= Step.
-  if (VF > 1 && Legal->requiresScalarEpilogue()) {
+  // Also, if the second last iteration val is extracted, we should follow this logic of
+  // (N - Step) iterations in vectorized loop and remaining (atmost Step)
+  // iterations in scalar loop.
+  if (VF > 1 &&
+      (Legal->requiresScalarEpilogue() || IsSecondLastIterationExtracted())) {
     auto *IsZero = Builder.CreateICmpEQ(R, ConstantInt::get(R->getType(), 0));
     R = Builder.CreateSelect(IsZero, Step, R);
   }
 
   VectorTripCount = Builder.CreateSub(TC, R, "n.vec");
 
   return VectorTripCount;
 }
void InnerLoopVectorizer::createEmptyLoop() {
 
   Value *StartIdx = ConstantInt::get(IdxTy, 0);
 
   // We need to test whether the backedge-taken count is uint##_max. Adding one
   // to it will cause overflow and an incorrect loop trip count in the vector
   // body. In case of overflow we want to directly jump to the scalar remainder
   // loop.
   emitMinimumIterationCountCheck(Lp, ScalarPH);
-  // Now, compare the new count to zero. If it is zero skip the vector loop and
-  // jump to the scalar loop.
+  // The loop step is equal to the vectorization factor (num of SIMD elements)
+  // times the unroll factor (num of SIMD instructions).
+  Value *CountRoundDown = getOrCreateVectorTripCount(Lp, OrigLoop);
   emitVectorLoopEnteredCheck(Lp, ScalarPH);
   // Generate the code to check any assumptions that we've made for SCEV
   // expressions.
   emitSCEVChecks(Lp, ScalarPH);
 
   // Generate the code that checks in runtime if arrays overlap. We put the
   // checks into a separate block to make the more common case of few elements
   // faster.
   emitMemRuntimeChecks(Lp, ScalarPH);
 
   // Generate the induction variable.
-  // The loop step is equal to the vectorization factor (num of SIMD elements)
-  // times the unroll factor (num of SIMD instructions).
-  Value *CountRoundDown = getOrCreateVectorTripCount(Lp);
   Constant *Step = ConstantInt::get(IdxTy, VF * UF);
   Induction =
       createInductionVariable(Lp, StartIdx, CountRoundDown, Step,
                               getDebugLocFromInstOrOperands(OldInduction));
 
   // We are going to resume the execution of the scalar loop.
   // Go over all of the induction variables that we found and fix the
   // PHIs that are left in the scalar version of the loop.
for (unsigned Part = 0; Part < UF; ++Part) {
     PhiParts[Part] = Shuffle;
     Incoming = PreviousParts[Part];
   }
 
   // Fix the latch value of the new recurrence in the vector loop.
   VecPhi->addIncoming(Incoming, LI->getLoopFor(LoopVectorBody)->getLoopLatch());
 
   // Extract the last vector element in the middle block. This will be the
-  // initial value for the recurrence when jumping to the scalar loop.
+  // initial value for the recurrence when jumping to the scalar loop.  Note
+  // that the last element need not always be the correct one, see
+  // IsSecondLastIterationExtracted. We fix the vector trip count so that the assumption
+  // of last element holds true. It would have been nice to fix this here (i.e.
+  // extract from the correct index), but extractElement does not allow anything
+  // other than a constant as the index.  We will not know if it is the last
+  // element or the element before the last, until runtime.
   auto *Extract = Incoming;
   if (VF > 1) {
     Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
     Extract = Builder.CreateExtractElement(Extract, Builder.getInt32(VF - 1),
                                            "vector.recur.extract");
   }
 
   // Fix the initial value of the original recurrence in the scalar loop.

test/Transforms/LoopVectorize/first-order-recurrence.ll

+; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -dce -S | FileCheck %s
+
+; This loop has 96 iterations, which divides evenly with the Step = 4. However, the vectorized part should only be upto 192.
+; The remaining scalar iterations will
+; be done at the remainder loop, so that we extract the correct value for
+; %val.phi = 94 + x.
+; If we vectorize the loop completely, we extract the last element of the vector val.phi (which would be incorrect: 95 + x) 
+; CHECK-LABEL: extract_second_last_iteration
+; CHECK: vector.body:
+; CHECK: %vec.ind = phi <4 x i32> [ <i32 0, i32 1, i32 2, i32 3>, %vector.ph ], [ %vec.ind.next, %vector.body ] 
+; CHECK: %[[L1:.+]] = add <4 x i32> %vec.ind, %broadcast.splat
+; CHECK: index.next = add i32 %index, 4
+; CHECK: icmp eq i32 %index.next, 92
+; CHECK: middle.block:
+; CHECK: %vector.recur.extract = extractelement <4 x i32> %[[L1]], i32 3
+; CHECK: scalar.ph:
+; CHECK: %scalar.recur.init = phi i32 [ %vector.recur.extract, %middle.block ], [ 0, %min.iters.checked ], [ 0, %entry ]
+; CHECK: %bc.resume.val = phi i32 [ 92, %middle.block ], [ 0, %entry ], [ 0, %min.iters.checked ] 
+; CHECK: for.body:
+; CHECK: %inc.phi = phi i32 [ %bc.resume.val, %scalar.ph ], [ %inc, %for.body ]
+; CHECK: %scalar.recur = phi i32 [ %scalar.recur.init, %scalar.ph ], [ %addx, %for.body ]
+; CHECK: %addx = add i32 %inc.phi, %x
+; CHECK: %cmp = icmp eq i32 %inc.phi, 95
+; CHECK: for.end:
+; CHECK: %val.phi.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %vector.recur.extract, %middle.block ]
+; CHECK: ret i32 %val.phi.lcssa
+define i32 @extract_second_last_iteration(i32* %cval, i32 %x)  {
+entry:
+  br label %for.body
+
+for.body:
+  %inc.phi = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %val.phi = phi i32 [ 0, %entry ], [ %addx, %for.body ]
+  %inc = add i32 %inc.phi, 1
+  %bc = zext i32 %inc.phi to i64
+  %addx = add i32 %inc.phi, %x
+  %cmp = icmp eq i32 %inc.phi, 95
+  br i1 %cmp, label %for.end, label %for.body
+
+for.end:
+  ret i32 %val.phi
+}
+
+; same as test above but we also store into index = val.phi.
+; FIXME: We could vectorize in this case following the same logic above (92
+; iterations vectorized and remaining is scalarized),
+; but currently the legality phase of the vectorizer prevents this.
+; CHECK-LABEL: storing_upto_second_last
+; CHECK-NOT: vector.body
+define i32 @storing_upto_second_last(i32* %cval, i32 %x)  {
+entry:
+  br label %for.body
+
+for.body:
+; val.phi at end of loop is 94 + x.
+  %inc.phi = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %val.phi = phi i32 [ 0, %entry ], [ %addx, %for.body ]
+  %inc = add nsw i32 %inc.phi, 1
+  %bc = zext i32 %val.phi to i64
+  %c40 = getelementptr inbounds i32, i32 * %cval, i64 %bc
+  %addx = add nuw nsw i32 %inc.phi, %x
+  store i32 %val.phi, i32* %c40
+  %cmp = icmp eq i32 %inc.phi, 95
+  br i1 %cmp, label %for.end, label %for.body
+
+for.end:
+  ret i32 0 
+}
+
+; We are extracting the last element from the loop. Make sure that we don't
+; unnecessarily create TC - Step iterations (i.e. 92 iterations) for the loop.
+; We can completely vectorize the loop at 96 iterations and extract the correct
+; value inc.phi = 95.
+; CHECK-LABEL: extract_last_iteration
+; CHECK: vector.body:
+; CHECK: %vec.ind = phi <4 x i32> [ <i32 0, i32 1, i32 2, i32 3>, %vector.ph ], [ %vec.ind.next, %vector.body ] 
+; CHECK: index.next = add i32 %index, 4
+; CHECK: icmp eq i32 %index.next, 96
+; CHECK: middle.block:
+; CHECK: %cmp.n = icmp eq i32 96, 96
+; CHECK: br i1 %cmp.n, label %for.end, label %scalar.ph
+; CHECK: for.end:
+; CHECK: %inc.phi.lcssa = phi i32 [ %inc.phi, %for.body ], [ 95, %middle.block ]
+; CHECK: ret i32 %inc.phi.lcssa
+define i32 @extract_last_iteration(i32* %cval, i32 %x)  {
+entry:
+  br label %for.body
+
+for.body:
+  %inc.phi = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %val.phi = phi i32 [ 0, %entry ], [ %addx, %for.body ]
+  %inc = add i32 %inc.phi, 1
+  %bc = zext i32 %inc.phi to i64
+  %addx = add i32 %inc.phi, %x
+  %cmp = icmp eq i32 %inc.phi, 95
+  br i1 %cmp, label %for.end, label %for.body
+
+for.end:
+  ret i32 %inc.phi
+}