[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 ↗

(On Diff #94877)

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 ↗

(On Diff #94877)

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.

llvm/trunk/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;
 
-  // 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 (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.
       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,
                                                       Type *Tp) {

llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp

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. 
+  // 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");
   }
 
   // Fix the initial value of the original recurrence in the scalar loop.

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

for.body:                                         ; preds = %entry, %for.body
   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: load <16 x i8>, <16 x i8>*
-; CHECK: add nuw nsw <16 x i32>
-; CHECK: store <16 x i8>
+; 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]
 define i8 @add_phifail2(i8* noalias nocapture readonly %p, i8* noalias nocapture %q, i32 %len) #0 {
 entry:
   br label %for.body
 
 for.cond.cleanup:                                 ; preds = %for.body, %entry
   %ret = trunc i32 %a_phi to i8
   ret i8 %ret
 

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

scalar.body:
   %tmp3 = add i64 0, 1
   %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
+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
+}