[LV] Fix the vector code generation for first order recurrence

[anna]

Summary

In first order recurrences where phi's are used outside the loop,
we should generate an additional vector.extract of the second last element from
the vectorized phi update.
This is because we require the phi itself (which is the value at the second last
iteration of the vector loop) and not the phi's update within the loop.

Fixes PR32396.

[anna]

Hi Matt,

Just FYI: The logic differs slightly from what you suggested in https://reviews.llvm.org/D31910#723721.
We need to extract from the correct index of the vectorized phi update. Extracting from index 3 of the vectorized phi (vector.recur), will give you the element at vector.update(3) - VF, which won't work.

middle.block:
  %vector.recur.extract = extractelement <4 x i32> %vector.update, i32 3
  %new.extract = extractelement <4 x i32> %vector.recur, i32 3

What we need is the the vectorized loop element equivalent to the second last element in the scalar loop (since the phi is the second last and the phi update is the last).

So, the logic used in this patch is:

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

[mssimpso]

In D31979#724852, @anna wrote:

Hi Matt,

Just FYI: The logic differs slightly from what you suggested in https://reviews.llvm.org/D31910#723721.
We need to extract from the correct index of the vectorized phi update. Extracting from index 3 of the vectorized phi (vector.recur), will give you the element at vector.update(3) - VF, which won't work.

OK, this makes sense to me.

lib/Transforms/Vectorize/LoopVectorize.cpp

4094–4095

We might as well create this extract unconditionally, regardless of whether the phi is used outside the loop. It would simplify the code here. So you could move it up to where "Extract" is created and not have to worry about setting "PhiHasUsersOutsideLoop". If there are no users, the extract is trivially dead and will be cleaned up by a later pass.

Unlike "Extract", I think we will need to initialize "ExtractForPhiUsedOutsideLoop" to something other than "Incoming" or nullptr. I think we'll need it when we are just unrolling and not vectorizing. Should it be PreviousParts[UF - 2]? See my comment in the test case.

It's probably a good idea to also rename "Extract" to something more descriptive since we now have two of them.

4113–4124

This looks a bit strange to me now. The LCSSA phi for "Phi" should always have an incoming value for "ExtractForPhiUsedOutsideLoop" and not "Extract". "Extract" will only be used to initialize the scalar recurrence. I just checked that we don't allow the update instruction to be used externally (I thought we could handle this), so we don't have to worry about that here.

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

371

Can you add a check for the VF = 1, UF > 1 case? I want to make sure the external use is correct when there is no extract.

[anna]

Need change based on Matt's comment to handle code gen when just unrolling and not vectorizing.

lib/Transforms/Vectorize/LoopVectorize.cpp

4094–4095

Agreed, will make the ExtractForPhiUsedOutsideLoop unconditional, and allow dce to eliminate it. I was under the impression these extracts would contribute to the cost model, but looking at the code, we only consider the instructions within the loop.

I'll take a look at the UF>1 and VF = 1 case. Thanks!

4113–4124

Ah! We can always use ExtractForPhiUsedOutsideLoop as the incoming value, since the loop is in LCSSA form. So, this Phi will be used only in LCSSA phis outside the loop, and it's occurrence will be only at the loop exit block.
Is that right?

How about these 2 checks under DEBUG:

1. Use the PhiHasUsersOutsideLoop to count number of users. It should be equal to one when we hit LCSSAPhi->getIncomingValue(0) == Phi. Or is this just checking LCSSA property? :)
2. update instruction has no external uses. If that's false, then we would get incorrect code gen here.

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

4113–4124

So, this Phi will be used only in LCSSA phis outside the loop, and it's occurrence will be only at the loop exit block. Is that right?

Yes, if any value in a loop is used externally, LCSSA will create a phi node for it in the exit block. The external uses will then be of the new phi.

How about these 2 checks under DEBUG:

For counting the external users, yes that would just be checking the LCSSA property, so you don't need to do that.

For the update instruction, you could do something like "assert(!hasOutsideLoopUser(Previous))". See hasOutsideLoopUser for more details.

You could also just go ahead and enable external uses of the update instruction (or do so in another patch if you prefer). The changes here to allow this would be to add "Extract" as the incoming value for the LCSSA phi for "Previous". You'll also need to add the update to AllowedExit over in canVectorizeInstrs.

[anna]

Handle code gen when loop unrolled but not vectorized.
Unconditionally generate the second extract, and other minor changes based on review.

[mssimpso]

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

360

Does the vector test actually pass? It doesn't look like it's well-formed. You need to define the regex for L1 and L2 when they're first used. Like %[[L1:.+]]. Also please indent the instructions under their corresponding block labels to make reading the test a little easier.

[anna]

Updated tests with correct prefix, fixed indentation.

[anna]

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

360

yeah, it actually did. These are silently ignored. We just need UNROLL-NO-IC without the CHECK suffix.

[mssimpso]

LGTM.

lib/Transforms/Utils/LoopUtils.cpp

if (Phi->getBasicBlockIndex(Preheader) < 0 ||
     return false;
 
   // Get the previous value. The previous value comes from the latch edge while
   // the initial value comes form the preheader edge.
   auto *Previous = dyn_cast<Instruction>(Phi->getIncomingValueForBlock(Latch));
   if (!Previous || !TheLoop->contains(Previous) || isa<PHINode>(Previous))
     return false;
 
-  for (User *U : Phi->users())
-    if (auto *I = dyn_cast<Instruction>(U)) {
-      // Ensure every user of the phi node is dominated by the previous value.
-      // The dominance requirement ensures the loop vectorizer will not need to
-      // vectorize the initial value prior to the first iteration of the loop.
+  for (User *U : Phi->users())
+    if (auto *I = dyn_cast<Instruction>(U)) {
       if (!DT->dominates(Previous, I))
         return false;
-      // When the phi node has users outside the loop, the current logic for
-      // fixFirstOrderRecurrences may generate incorrect code. Specifically, we
-      // extract the last element from the vectorized phi, which would be the
-      // update to the phi before exiting the loop. However, what we want is the
-      // previous phi value before the update (i.e. the second last update
-      // before end of the vectorized loop).
-      // See added test cases in first-order-recurrence.ll
-      if (!TheLoop->contains(I))
-        return false;
     }
 
   return true;
 }
 
 /// This function returns the identity element (or neutral element) for
 /// the operation K.
 Constant *RecurrenceDescriptor::getRecurrenceIdentity(RecurrenceKind K,

lib/Transforms/Vectorize/LoopVectorize.cpp

for (unsigned Part = 0; Part < UF; ++Part) {
     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. 
-  // FIXME: Note that the last vector element need not always be the correct one:
-  // consider a loop  where we have phi uses outside the loop - we need the
-  // second last iteration value and not the last one). For now, we avoid
-  // considering such cases as firstOrderRecurrences (see
-  // isFirstOrderRecurrence).
   auto *Extract = Incoming;
   if (VF > 1) {
     Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
     Extract = Builder.CreateExtractElement(Extract, Builder.getInt32(VF - 1),
                                            "vector.recur.extract");
   }
+  // We also need to extract the second last element in the middle block if the
+  // Phi is used outside the loop. We need to extract the phi itself
+  // and not the last element (the phi update in the current iteration). This
+  // will be the value when jumping to the exit block from the LoopMiddleBlock,
+  // when the scalar loop is not run at all.
+  auto PhiHasUsersOutsideLoop = false;
+  for (User *U : Phi->users()) {
+    if (auto *I = dyn_cast<Instruction>(U))
+      if (!OrigLoop->contains(I)) {
+        PhiHasUsersOutsideLoop = true;
+        break;
+      }
+  }
+  Value *ExtractForPhiUsedOutsideLoop = nullptr;
+  if (VF > 1 && PhiHasUsersOutsideLoop)
+    ExtractForPhiUsedOutsideLoop = Builder.CreateExtractElement(
+        Incoming, Builder.getInt32(VF - 2), "vector.recur.extract.for.phi");
 
   // Fix the initial value of the original recurrence in the scalar loop.
   Builder.SetInsertPoint(&*LoopScalarPreHeader->begin());
   auto *Start = Builder.CreatePHI(Phi->getType(), 2, "scalar.recur.init");
   for (auto *BB : predecessors(LoopScalarPreHeader)) {
     auto *Incoming = BB == LoopMiddleBlock ? Extract : ScalarInit;
     Start->addIncoming(Incoming, BB);
   }
 
   Phi->setIncomingValue(Phi->getBasicBlockIndex(LoopScalarPreHeader), Start);
   Phi->setName("scalar.recur");
 
   // Finally, fix users of the recurrence outside the loop. The users will need
   // either the last value of the scalar recurrence or the last value of the
   // vector recurrence we extracted in the middle block. Since the loop is in
   // LCSSA form, we just need to find the phi node for the original scalar
   // recurrence in the exit block, and then add an edge for the middle block.
+  // Also, choose the correct value of the vector recurrence for fixing up users of
+  // the recurrence outside the loop.
+  if (PhiHasUsersOutsideLoop)
+    Extract = ExtractForPhiUsedOutsideLoop;
   for (auto &I : *LoopExitBlock) {
     auto *LCSSAPhi = dyn_cast<PHINode>(&I);
     if (!LCSSAPhi)
       break;
     if (LCSSAPhi->getIncomingValue(0) == Phi) {
       LCSSAPhi->addIncoming(Extract, LoopMiddleBlock);
       break;
     }
   }
 }
 
 void InnerLoopVectorizer::fixReduction(PHINode *Phi) {
   Constant *Zero = Builder.getInt32(0);
 
   // Get it's reduction variable descriptor.
   assert(Legal->isReductionVariable(Phi) &&

test/Transforms/LoopVectorize/AArch64/loop-vectorization-factors.ll

for.body:                                         ; preds = %entry, %for.body
   %arrayidx3 = getelementptr inbounds i8, i8* %q, i64 %indvars.iv
   store i8 %conv1, i8* %arrayidx3
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %lftr.wideiv = trunc i64 %indvars.iv.next to i32
   %exitcond = icmp eq i32 %lftr.wideiv, %len
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 }
 
-; CHECK-LABEL: @add_phifail2(
-; CHECK-NOT: load <16 x i8>, <16 x i8>*
-; CHECK-NOT: add nuw nsw <16 x i32>
-; CHECK-NOT: store <16 x i8>
 ; Function Attrs: nounwind
-; FIXME: Currently, if we vectorize this loop, we will generate incorrect code
-; if %len evenly divides VF. Vectorized loop code gen returns a_phi = p[len -1],
-; whereas it should be the previous value a_phi = p[len -2]
+; When we vectorize this loop, we generate correct code
+; even when %len exactly divides VF (since we extract from the second last index
+; and pass this to the for.cond.cleanup block). Vectorized loop returns 
+; the correct value a_phi = p[len -2]
 define i8 @add_phifail2(i8* noalias nocapture readonly %p, i8* noalias nocapture %q, i32 %len) #0 {
+; CHECK-LABEL: @add_phifail2(
+; CHECK: vector.body:
+; CHECK: %wide.load = load <16 x i8>, <16 x i8>*
+; CHECK: %[[L1:.+]] = zext <16 x i8> %wide.load to <16 x i32>
+; CHECK: add nuw nsw <16 x i32>
+; CHECK: store <16 x i8>
+; CHECK: add i64 %index, 16
+; CHECK: icmp eq i64 %index.next, %n.vec
+; CHECK: middle.block:
+; CHECK: %vector.recur.extract = extractelement <16 x i32> %[[L1]], i32 15
+; CHECK: %vector.recur.extract.for.phi = extractelement <16 x i32> %[[L1]], i32 14
+; CHECK: for.cond.cleanup:
+; CHECK: %a_phi.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %vector.recur.extract.for.phi, %middle.block ]
+; CHECK: %ret = trunc i32 %a_phi.lcssa to i8
+; CHECK: ret i8 %ret
 entry:
   br label %for.body
 
 for.cond.cleanup:                                 ; preds = %for.body, %entry
   %ret = trunc i32 %a_phi to i8
   ret i8 %ret
 
 for.body:                                         ; preds = %entry, %for.body

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

scalar.body:
   %i.next = add nuw nsw i64 %i, 1
   %cond = icmp eq i64 %i.next, undef
   br i1 %cond, label %for.end, label %scalar.body
 
 for.end:
   ret void
 }
 
-; FIXME: we can vectorize this first order recurrence, by generating two
-; extracts - one for the phi `val.phi` and other for the phi update `addx`.
-; val.phi at end of loop is 94 + x.
-; CHECK-LABEL: extract_second_last_iteration
-; CHECK-NOT: vector.body
+; We vectorize this first order recurrence, by generating two
+; extracts for the phi `val.phi` - one at the last index and 
+; another at the second last index. We need these 2 extracts because 
+; the first order recurrence phi is used outside the loop, so we require the phi
+; itself and not its update (addx).
+; UNROLL-NO-IC-LABEL: extract_second_last_iteration
+; UNROLL-NO-IC-CHECK: vector.body
+; UNROLL-NO-IC-CHECK: %vector.recur = phi <4 x i32> [ <i32 undef, i32 undef, i32 undef, i32 0>, %vector.ph ], [ %[[L1], %vector.body ]
+; UNROLL-NO-IC-CHECK: %step.add = add <4 x i32> %vec.ind, <i32 4, i32 4, i32 4, i32 4>
+; UNROLL-NO-IC-CHECK: %[[L2]] = add <4 x i32> %vec.ind, %broadcast.splat
+; UNROLL-NO-IC-CHECK: %[[L1]] = add <4 x i32> %step.add, %broadcast.splat
+; UNROLL-NO-IC-CHECK: %index.next = add i32 %index, 8
+; UNROLL-NO-IC-CHECK: icmp eq i32 %index.next, 96
+; UNROLL-NO-IC-CHECK: middle.block
+; UNROLL-NO-IC-CHECK: icmp eq i32 96, 96
+; UNROLL-NO-IC-CHECK: %vector.recur.extract = extractelement <4 x i32> %[[L1], i32 3
+; UNROLL-NO-IC-CHECK: %vector.recur.extract.for.phi = extractelement <4 x i32> %[[L1], i32 2
+; UNROLL-NO-IC-CHECK: for.end
+; UNROLL-NO-IC-CHECK: %val.phi.lcssa = phi i32 [ %scalar.recur, %for.body ], [ %vector.recur.extract.for.phi, %middle.block ]
 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
 }