[LV] Generalize conditions for sinking instrs for first order recurrences.

[fhahn]

Summary

If the recurrence PHI node has a single user, we can sink any
instruction without side effects, given that all users are dominated by
the instruction computing the incoming value of the next iteration
('Previous'). We can sink instructions that may cause traps, because
that only causes the trap to occur later, but not on any new paths.

With the relaxed check, we also have to make sure that we do not have a
direct cycle (meaning PHI user == 'Previous), which indicates a
reduction relation, which potentially gets missed by
ReductionDescriptor.

As follow-ups, we can also sink stores, iff they do not alias with
other instructions we move them across and we could also support sinking
chains of instructions and multiple users of the PHI.

Fixes PR43398.

[Ayal]

llvm/lib/Analysis/IVDescriptors.cpp

702–711

Above TODO comment should be updated as well (thanks :-). It's possible to further extend sinking loads/stores as noted, which will affect interleave-group analysis; and also to sinking multiple users of Phi, which will need to preserve relative order to maintain dependences between themselves.

705

by the instruction[s]

706

canSink >> canSinkAfter ?

auto >> bool ?

712–713

Maybe better to first check if Previous already dominates I, before checking all of I's uses.

715

That's true, but it isn't checked below if Phi has more than one use, and it isn't checked in canSink (if any of I's uses is equal to SinkPoint). Perhaps it should be an assert, in all these places, verifying the precondition that Phi does not belong to any reduction cycle.

729–730

The above for loop looks similar to the canSink[After]() lambda; maybe rename it allUsesDominatedBy() or something, and reapply here?

[fhahn]

Address review comments, thanks Ayal!

[fhahn]

llvm/lib/Analysis/IVDescriptors.cpp

706

I used auto for the lambda, but I've got rid of the lambda for now and moved it down.

715

Agreed. I guess to be sure we would have to recursively check all users, until we either reach a cycle or are outside the loop, right? I've added the (simple) check here because this is very easy to run into in practice. Would it be OK to have the assertion as follow-up?

[Ayal]

This looks good to me, plus some final minor comments, thanks!

llvm/lib/Analysis/IVDescriptors.cpp

715

My bad, the added check is correct; cannot_sink_reduction() test below demonstrates how this can now happen. (It could not happen until now because a cast must change types :-). And dominates(DominatedBy, U) fails if DominatedBy==U, so the users of Phi below are covered.

"and there's nothing to sink" >> "which cannot be handled by sinking" (an instruction cannot sink after itself)

728

It may be better to check allUsesDominatedBy() last, i.e., after checking for same Parents, in terms of compile-time.

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

7

Worth commenting that the test is derived from PR43398, along with a simple version in C.

63

Note that a trapping instruction may sink past a trapping instruction; but if loop is vectorized, the order may swap anyhow...

122

Alternatively CHECK-NOT that no vector type got generated? Holds for additional not-to-be-vectorized tests below.

163

no[t] alias

[fhahn]

Thanks Ayal, committed with the comments addressed.

llvm/lib/Analysis/IVDescriptors.cpp

728

I've moved the checks in the if directly, no need to trust the compiler to sink the checks.

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

122

I've kept the more verbose check lines, as it seems a bit more explicit (and they are still quite compact IMO)

[Ayal]

llvm/lib/Analysis/IVDescriptors.cpp

732

Post-commit note, to be considered e.g. when extending sinkings further: seems more logical to start with "if (allUsesDominatedBy(Phi, Previous)) return true; // We already are good w/o sinking", then try to sink with "if (Phi->hasOneUse()) {...}" and finally return false if sinking failed.

[Ayal]

As revert eaff300 shows, more care must be taken to avoid having one FOR (First Order Recurring) phi sink an instruction that another FOR phi needs to keep in place. Namely, any Previous instruction that feeds a FOR phi across the backarc must not be allowed to sink, because that may break its dominance on other instructions who use the FOR phi. The original check against SinkAfter.count(Previous)) // Cannot rely on dominance due to motion addresses this concern but only partially; it works if the FOR phi that needs to sink an instruction is handled before the FOR phi that needs to keep it in place. But the two phi's could be encountered in the other order.

One way to fix this is to record every "Previous" instruction using the existing SinkAfter map, as illustrated below. Another way may be to post-process all FirstOrderRecurrences along with SinkAfter.

diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp
index bdebd71d7f6..fd78a8335e0 100644
--- a/llvm/lib/Analysis/IVDescriptors.cpp
+++ b/llvm/lib/Analysis/IVDescriptors.cpp
@@ -696,6 +696,9 @@ bool RecurrenceDescriptor::isFirstOrderRecurrence(
       SinkAfter.count(Previous)) // Cannot rely on dominance due to motion.
     return false;
 
+  // Record that Previous must stay in place, as it must continue to dominate other instructions.
+  SinkAfter[Previous] = nullptr;
+
   // 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.
@@ -717,6 +720,11 @@ bool RecurrenceDescriptor::isFirstOrderRecurrence(
     if (Previous == I)
       return false;
 
+    if (SinkAfter.count(I)) {
+      LLVM_DEBUG(dbgs() << "LV: cannot sink instruction " << *I << " because it is involved in multiple recurrence phi's.\n");
+      return false;
+    }
+
     if (DT->dominates(Previous, I)) // We already are good w/o sinking.
       return true;
 
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 44623dd0ebc..5ca40dc3974 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7152,7 +7152,7 @@ VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
       // handling this instruction until after we've handled the instruction it
       // should follow.
       auto SAIt = SinkAfter.find(Instr);
-      if (SAIt != SinkAfter.end()) {
+      if (SAIt != SinkAfter.end() && SAIt->second) {
         LLVM_DEBUG(dbgs() << "Sinking" << *SAIt->first << " after"
                           << *SAIt->second
                           << " to vectorize a 1st order recurrence.\n");

[fhahn]

Thanks Ayal, I’ll add a fix and recommit the patch soon

[fhahn]

In D69228#1737857, @Ayal wrote:

As revert eaff300 shows, more care must be taken to avoid having one FOR (First Order Recurring) phi sink an instruction that another FOR phi needs to keep in place. Namely, any Previous instruction that feeds a FOR phi across the backarc must not be allowed to sink, because that may break its dominance on other instructions who use the FOR phi. The original check against SinkAfter.count(Previous)) // Cannot rely on dominance due to motion addresses this concern but only partially; it works if the FOR phi that needs to sink an instruction is handled before the FOR phi that needs to keep it in place. But the two phi's could be encountered in the other order.

I've recommitted the patch: https://github.com/llvm/llvm-project/commit/9d24933f79dd7db3e469b3c4402e076cc41418f7.

It checks for all found recurrence PHIs, if any of the previous values needs sinking and bails out in LoopVectorizationLegality. I think there's still a bit of cleanup potential, e.g. removing another check from isFirstOrderRecurrence and maybe moving it to a function (e.g. validateFirstOrderRecurrences in IVDescriptors). I think we could also handle additional cases, where the required dominance relationships still hold after sinking. Let's do that as follow ups, once the patch sticks.

llvm/lib/Analysis/IVDescriptors.cpp

bool RecurrenceDescriptor::isFirstOrderRecurrence(
   auto *Previous = dyn_cast<Instruction>(Phi->getIncomingValueForBlock(Latch));
   if (!Previous || !TheLoop->contains(Previous) || isa<PHINode>(Previous) ||
       SinkAfter.count(Previous)) // Cannot rely on dominance due to motion.
     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.
-  // TODO: Consider extending this sinking to handle other kinds of instructions
-  // and expressions, beyond sinking a single cast past Previous.
+  // TODO: Consider extending this sinking to handle memory instructions and
+  // phis with multiple users.
+
+  // Returns true, if all users of I are dominated by DominatedBy.
+  auto allUsesDominatedBy = [DT](Instruction *I, Instruction *DominatedBy) {
+    return all_of(I->uses(), [DT, DominatedBy](Use &U) {
+      return DT->dominates(DominatedBy, U);
+    });
+  };
+
   if (Phi->hasOneUse()) {
-    auto *I = Phi->user_back();
-    if (I->isCast() && (I->getParent() == Phi->getParent()) && I->hasOneUse() &&
-        DT->dominates(Previous, I->user_back())) {
-      if (!DT->dominates(Previous, I)) // Otherwise we're good w/o sinking.
-        SinkAfter[I] = Previous;
+    Instruction *I = Phi->user_back();
+
+    // If the user of the PHI is also the incoming value, we potentially have a
+    // reduction and which cannot be handled by sinking.
+    if (Previous == I)
+      return false;
+
+    if (DT->dominates(Previous, I)) // We already are good w/o sinking.
+      return true;
+
+    // We can sink any instruction without side effects, as long as all users
+    // are dominated by the instruction we are sinking after.
+    if (I->getParent() == Phi->getParent() && !I->mayHaveSideEffects() &&
+        allUsesDominatedBy(I, Previous)) {
       return true;
     }
   }
 
-  for (User *U : Phi->users())
-    if (auto *I = dyn_cast<Instruction>(U)) {
-      if (!DT->dominates(Previous, I))
-        return false;
-    }
-
-  return true;
+  return allUsesDominatedBy(Phi, Previous);
 }
 
 /// This function returns the identity element (or neutral element) for
 /// the operation K.
 Constant *RecurrenceDescriptor::getRecurrenceIdentity(RecurrenceKind K,
                                                       Type *Tp) {
   switch (K) {
   case RK_IntegerXor:

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

+; RUN: opt -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S %s | FileCheck %s
+
+
+@p = external local_unnamed_addr global [257 x i32], align 16
+@q = external local_unnamed_addr global [257 x i32], align 16
+
+; Test case for PR43398.
+
+define void @can_sink_after_store(i32 %x, i32* %ptr, i64 %tc) local_unnamed_addr #0 {
+; CHECK-LABEL: vector.ph:
+; CHECK:        %broadcast.splatinsert1 = insertelement <4 x i32> undef, i32 %x, i32 0
+; CHECK-NEXT:   %broadcast.splat2 = shufflevector <4 x i32> %broadcast.splatinsert1, <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:   %vector.recur.init = insertelement <4 x i32> undef, i32 %.pre, i32 3
+; CHECK-NEXT:    br label %vector.body
+
+; CHECK-LABEL: vector.body:
+; CHECK-NEXT:   %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+; CHECK-NEXT:   %vector.recur = phi <4 x i32> [ %vector.recur.init, %vector.ph ], [ %wide.load, %vector.body ]
+; CHECK-NEXT:   %offset.idx = add i64 1, %index
+; CHECK-NEXT:   %broadcast.splatinsert = insertelement <4 x i64> undef, i64 %offset.idx, i32 0
+; CHECK-NEXT:   %broadcast.splat = shufflevector <4 x i64> %broadcast.splatinsert, <4 x i64> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:   %induction = add <4 x i64> %broadcast.splat, <i64 0, i64 1, i64 2, i64 3>
+; CHECK-NEXT:   %0 = add i64 %offset.idx, 0
+; CHECK-NEXT:   %1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %0
+; CHECK-NEXT:   %2 = getelementptr inbounds i32, i32* %1, i32 0
+; CHECK-NEXT:   %3 = bitcast i32* %2 to <4 x i32>*
+; CHECK-NEXT:   %wide.load = load <4 x i32>, <4 x i32>* %3, align 4
+; CHECK-NEXT:   %4 = shufflevector <4 x i32> %vector.recur, <4 x i32> %wide.load, <4 x i32> <i32 3, i32 4, i32 5, i32 6>
+; CHECK-NEXT:   %5 = add <4 x i32> %4, %broadcast.splat2
+; CHECK-NEXT:   %6 = add <4 x i32> %5, %wide.load
+; CHECK-NEXT:   %7 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %0
+; CHECK-NEXT:   %8 = getelementptr inbounds i32, i32* %7, i32 0
+; CHECK-NEXT:   %9 = bitcast i32* %8 to <4 x i32>*
+; CHECK-NEXT:   store <4 x i32> %6, <4 x i32>* %9, align 4
+; CHECK-NEXT:   %index.next = add i64 %index, 4
+; CHECK-NEXT:   %10 = icmp eq i64 %index.next, 1996
+; CHECK-NEXT:   br i1 %10, label %middle.block, label %vector.body
+;
+entry:
+  br label %preheader
+
+preheader:
+  %idx.phi.trans = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 1
+  %.pre = load i32, i32* %idx.phi.trans, align 4
+  br label %for
+
+for:
+  %pre.phi = phi i32 [ %.pre, %preheader ], [ %pre.next, %for ]
+  %iv = phi i64 [ 1, %preheader ], [ %iv.next, %for ]
+  %add.1 = add i32 %pre.phi, %x
+  %idx.1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %iv
+  %pre.next = load i32, i32* %idx.1, align 4
+  %add.2 = add i32 %add.1, %pre.next
+  %idx.2 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %iv
+  store i32 %add.2, i32* %idx.2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 2000
+  br i1 %exitcond, label %exit, label %for
+
+exit:
+  ret void
+}
+
+; We can sink potential trapping instructions, as this will only delay the trap
+; and not introduce traps on additional paths.
+define void @sink_sdiv(i32 %x, i32* %ptr, i64 %tc) local_unnamed_addr #0 {
+; CHECK-LABEL: vector.ph:
+; CHECK:        %broadcast.splatinsert1 = insertelement <4 x i32> undef, i32 %x, i32 0
+; CHECK-NEXT:   %broadcast.splat2 = shufflevector <4 x i32> %broadcast.splatinsert1, <4 x i32> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:   %vector.recur.init = insertelement <4 x i32> undef, i32 %.pre, i32 3
+; CHECK-NEXT:    br label %vector.body
+
+; CHECK-LABEL: vector.body:
+; CHECK-NEXT:   %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+; CHECK-NEXT:   %vector.recur = phi <4 x i32> [ %vector.recur.init, %vector.ph ], [ %wide.load, %vector.body ]
+; CHECK-NEXT:   %offset.idx = add i64 1, %index
+; CHECK-NEXT:   %broadcast.splatinsert = insertelement <4 x i64> undef, i64 %offset.idx, i32 0
+; CHECK-NEXT:   %broadcast.splat = shufflevector <4 x i64> %broadcast.splatinsert, <4 x i64> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT:   %induction = add <4 x i64> %broadcast.splat, <i64 0, i64 1, i64 2, i64 3>
+; CHECK-NEXT:   %0 = add i64 %offset.idx, 0
+; CHECK-NEXT:   %1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %0
+; CHECK-NEXT:   %2 = getelementptr inbounds i32, i32* %1, i32 0
+; CHECK-NEXT:   %3 = bitcast i32* %2 to <4 x i32>*
+; CHECK-NEXT:   %wide.load = load <4 x i32>, <4 x i32>* %3, align 4
+; CHECK-NEXT:   %4 = shufflevector <4 x i32> %vector.recur, <4 x i32> %wide.load, <4 x i32> <i32 3, i32 4, i32 5, i32 6>
+; CHECK-NEXT:   %5 = sdiv <4 x i32> %4, %broadcast.splat2
+; CHECK-NEXT:   %6 = add <4 x i32> %5, %wide.load
+; CHECK-NEXT:   %7 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %0
+; CHECK-NEXT:   %8 = getelementptr inbounds i32, i32* %7, i32 0
+; CHECK-NEXT:   %9 = bitcast i32* %8 to <4 x i32>*
+; CHECK-NEXT:   store <4 x i32> %6, <4 x i32>* %9, align 4
+; CHECK-NEXT:   %index.next = add i64 %index, 4
+; CHECK-NEXT:   %10 = icmp eq i64 %index.next, 1996
+; CHECK-NEXT:   br i1 %10, label %middle.block, label %vector.body
+;
+entry:
+  br label %preheader
+
+preheader:
+  %idx.phi.trans = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 1
+  %.pre = load i32, i32* %idx.phi.trans, align 4
+  br label %for
+
+for:
+  %pre.phi = phi i32 [ %.pre, %preheader ], [ %pre.next, %for ]
+  %iv = phi i64 [ 1, %preheader ], [ %iv.next, %for ]
+  %div.1 = sdiv i32 %pre.phi, %x
+  %idx.1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %iv
+  %pre.next = load i32, i32* %idx.1, align 4
+  %add.2 = add i32 %div.1, %pre.next
+  %idx.2 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %iv
+  store i32 %add.2, i32* %idx.2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 2000
+  br i1 %exitcond, label %exit, label %for
+
+exit:
+  ret void
+}
+
+; FIXME: Currently we can only sink a single instruction. For the example below,
+;        we also have to sink users.
+define void @cannot_sink_with_additional_user(i32 %x, i32* %ptr, i64 %tc) {
+; CHECK-LABEL: define void @cannot_sink_with_additional_user(
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   br label %preheader
+
+; CHECK-LABEL: preheader:                                        ; preds = %entry
+; CHECK:  br label %for
+
+; CHECK-LABEL: for:                                              ; preds = %for, %preheader
+; CHECK  br i1 %exitcond, label %exit, label %for
+
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    ret void
+
+entry:
+  br label %preheader
+
+preheader:
+  %idx.phi.trans = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 1
+  %.pre = load i32, i32* %idx.phi.trans, align 4
+  br label %for
+
+for:
+  %pre.phi = phi i32 [ %.pre, %preheader ], [ %pre.next, %for ]
+  %iv = phi i64 [ 1, %preheader ], [ %iv.next, %for ]
+  %add.1 = add i32 %pre.phi, %x
+  %add.2 = add i32 %add.1, %x
+  %idx.1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %iv
+  %pre.next = load i32, i32* %idx.1, align 4
+  %add.3 = add i32 %add.1, %pre.next
+  %add.4 = add i32 %add.2, %add.3
+  %idx.2 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %iv
+  store i32 %add.4, i32* %idx.2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 2000
+  br i1 %exitcond, label %exit, label %for
+
+exit:
+  ret void
+}
+
+; FIXME: We can sink a store, if we can guarantee that it does not alias any
+;        loads/stores in between.
+define void @cannot_sink_store(i32 %x, i32* %ptr, i64 %tc) {
+; CHECK-LABEL: define void @cannot_sink_store(
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   br label %preheader
+
+; CHECK-LABEL: preheader:                                        ; preds = %entry
+; CHECK:  br label %for
+
+; CHECK-LABEL: for:                                              ; preds = %for, %preheader
+; CHECK  br i1 %exitcond, label %exit, label %for
+
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %preheader
+
+preheader:
+  %idx.phi.trans = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 1
+  %.pre = load i32, i32* %idx.phi.trans, align 4
+  br label %for
+
+for:
+  %pre.phi = phi i32 [ %.pre, %preheader ], [ %pre.next, %for ]
+  %iv = phi i64 [ 1, %preheader ], [ %iv.next, %for ]
+  %add.1 = add i32 %pre.phi, %x
+  store i32 %add.1, i32* %ptr
+  %idx.1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %iv
+  %pre.next = load i32, i32* %idx.1, align 4
+  %add.2 = add i32 %add.1, %pre.next
+  %idx.2 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %iv
+  store i32 %add.2, i32* %idx.2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 2000
+  br i1 %exitcond, label %exit, label %for
+
+exit:
+  ret void
+}
+
+; Some kinds of reductions are not detected by IVDescriptors. If we have a
+; cycle, we cannot sink it.
+define void @cannot_sink_reduction(i32 %x, i32* %ptr, i64 %tc) {
+; CHECK-LABEL: define void @cannot_sink_reduction(
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   br label %preheader
+
+; CHECK-LABEL: preheader:                                        ; preds = %entry
+; CHECK:  br label %for
+
+; CHECK-LABEL: for:                                              ; preds = %for, %preheader
+; CHECK  br i1 %exitcond, label %exit, label %for
+
+; CHECK-LABEL: exit:                                    ; preds = %for
+; CHECK-NET:     ret void
+;
+entry:
+  br label %preheader
+
+preheader:
+  %idx.phi.trans = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 1
+  %.pre = load i32, i32* %idx.phi.trans, align 4
+  br label %for
+
+for:
+  %pre.phi = phi i32 [ %.pre, %preheader ], [ %d, %for ]
+  %iv = phi i64 [ 1, %preheader ], [ %iv.next, %for ]
+  %d = sdiv i32 %pre.phi, %x
+  %idx.1 = getelementptr inbounds [257 x i32], [257 x i32]* @p, i64 0, i64 %iv
+  %pre.next = load i32, i32* %idx.1, align 4
+  %add.2 = add i32 %x, %pre.next
+  %idx.2 = getelementptr inbounds [257 x i32], [257 x i32]* @q, i64 0, i64 %iv
+  store i32 %add.2, i32* %idx.2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 2000
+  br i1 %exitcond, label %exit, label %for
+
+exit:
+  ret void
+}