[LV] Sink casts to unravel first order recurrence

[Ayal]

Summary

A loop-phi Phi that involves a first order recurrence can be vectorized only if the Previous instruction that feeds Phi from the previous iteration dominates all of Phi's users. This condition is needed to allow introducing instructions that use both Phi and Previous, combine their values, and feed the users of Phi and Previous.

When Phi has users of types different from its own, a mediating Cast instruction that uses Phi may appear above Previous, thereby breaking this condition. This patch tries to sink such a Cast past Previous, provided the users of Cast are dominated by Previous, thereby enabling vectorization of the loop.

[mssimpso]

Hi Ayal,

Thanks for extending this. The current dominance requirement can be fairly restrictive. Is there anything special about casts? I'm wondering if it's worth extending this to handle other kinds of instructions/expressions.

[aemerson]

In general I wonder if this is really the best place to do this. It would be nice if the loop was canonicalised to be in this form given how cheap it is to do. Perhaps LoopSimplify? Not blocking this change, but something to think about.

lib/Transforms/Utils/LoopUtils.cpp

556 ↗

(On Diff #98502)

In what situation can this occur? Isn't this an output vector?

564 ↗

(On Diff #98502)

Bug here, bitwise & used. Which is also telling me that there aren't sufficient tests for this, especially the case where Previous doesn't dominate the cast's user and so the cast can't be legally sunk past Previous.

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

401 ↗

(On Diff #98502)

Can we add a few more checks to verify the result of the vectorization is sane?

[Ayal]

In D33058#751569, @mssimpso wrote:

Hi Ayal,

Thanks for extending this. The current dominance requirement can be fairly restrictive. Is there anything special about casts? I'm wondering if it's worth extending this to handle other kinds of instructions/expressions.

The case of casts arises naturally when GVN::PerformLoadPRE() optimizes an a[i]+a[i+1] where "a" is an array e.g. of type short and the addition is of type int, as in the attached test-case. It is also possible though to apply this optimization manually to the source code, by doing:

short ai = a[0];
for (int i = 0; i < n; ++i) {
  short aiPlusOne = a[i+1];
  b[i] = ai * aiPlusOne;
  ai = aiPlusOne;
}

or manually add any instruction to use Phi before Previous, such as the following "+= 3":

int ai = a[0];
for (int i = 0; i < n; ++i) {
  ai += 3;
  int aiPlusOne = a[i+1];
  b[i] = ai * aiPlusOne;
  ai = aiPlusOne;
}

(Doing "b[i]=(a[i]+3)+a[i+1]" btw produces a cast followed by an add, both above Previous; a test-case @aemerson asked for.)

In D33058#752015, @aemerson wrote:

In general I wonder if this is really the best place to do this. It would be nice if the loop was canonicalised to be in this form given how cheap it is to do. Perhaps LoopSimplify? Not blocking this change, but something to think about.

The specific a[i]+a[i+1] case could indeed be handled before vectorization, namely by (LoopSimplify?) hoisting the cast along with the load. This applies in general to other instructions that operate symmetrically on both a[i] and a[i+1]. Instructions that only apply to a[i] need to be carefully placed inside the loop in order for it to be vectorized. It's probably better to have the vectorizer handle such motion; nothing keeps these instructions from returning back to their original positions.

Sounds reasonable?

lib/Transforms/Utils/LoopUtils.cpp

556 ↗

(On Diff #98502)

In general, once instructions are allowed to move, their updated location and dominance info should be considered. Furthermore, if transitive motion is allowed, e.g., sink a after b and sink b after c, they need to execute in the right order.

In the scenarios encountered Previous is a load, so this does not occur, as sinking is restricted here to casts.

Note that Interleave Groups also move instructions, but w/o affecting such sinking of casts: loads move upwards (w/o breaking their original dominance info) and stores move downwards (having no users).

564 ↗

(On Diff #98502)

Right, thanks!

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

401 ↗

(On Diff #98502)

Sure.

[aemerson]

In D33058#752230, @Ayal wrote:

The specific a[i]+a[i+1] case could indeed be handled before vectorization, namely by (LoopSimplify?) hoisting the cast along with the load. This applies in general to other instructions that operate symmetrically on both a[i] and a[i+1]. Instructions that only apply to a[i] need to be carefully placed inside the loop in order for it to be vectorized. It's probably better to have the vectorizer handle such motion; nothing keeps these instructions from returning back to their original positions.

Sounds reasonable?

Yep, thanks.

[Ayal]

In D33058#751569, @mssimpso wrote:
I'm wondering if it's worth extending this to handle other kinds of instructions/expressions.

Updated version captures this wondering in a TODO, along with addressing the other review comments.

[Ayal]

lib/Transforms/Utils/LoopUtils.cpp

564 ↗

(On Diff #98502)

Fixed. (Missing '&' fell off column 80 I guess ;-)

Added a test as mentioned above where Previous (the load) doesn't dominate the cast's user, so the cast cannot be legally sunk past Previous w/o sinking this user along with it.

[mssimpso]

LGTM, assuming Amara has no more comments. Thanks Ayal!

[aemerson]

Sorry, this fell off my radar. LGTM.

llvm/trunk/include/llvm/Transforms/Utils/LoopUtils.h

public:
   /// Returns true if Phi is a reduction in TheLoop. The RecurrenceDescriptor is
   /// returned in RedDes.
   static bool isReductionPHI(PHINode *Phi, Loop *TheLoop,
                              RecurrenceDescriptor &RedDes);
 
   /// Returns true if Phi is a first-order recurrence. A first-order recurrence
   /// is a non-reduction recurrence relation in which the value of the
   /// recurrence in the current loop iteration equals a value defined in the
-  /// previous iteration.
-  static bool isFirstOrderRecurrence(PHINode *Phi, Loop *TheLoop,
-                                     DominatorTree *DT);
+  /// previous iteration. \p SinkAfter includes pairs of instructions where the
+  /// first will be rescheduled to appear after the second if/when the loop is
+  /// vectorized. It may be augmented with additional pairs if needed in order
+  /// to handle Phi as a first-order recurrence.
+  static bool
+  isFirstOrderRecurrence(PHINode *Phi, Loop *TheLoop,
+                         DenseMap<Instruction *, Instruction *> &SinkAfter,
 
   RecurrenceKind getRecurrenceKind() { return Kind; }
 
   MinMaxRecurrenceKind getMinMaxRecurrenceKind() { return MinMaxKind; }
 
   TrackingVH<Value> getRecurrenceStartValue() { return StartValue; }
 
   Instruction *getLoopExitInstr() { return LoopExitInstr; }

llvm/trunk/lib/Transforms/Utils/LoopUtils.cpp

bool RecurrenceDescriptor::isReductionPHI(PHINode *Phi, Loop *TheLoop,
   if (AddReductionVar(Phi, RK_FloatMinMax, TheLoop, HasFunNoNaNAttr, RedDes)) {
     DEBUG(dbgs() << "Found an float MINMAX reduction PHI." << *Phi << "\n");
     return true;
   }
   // Not a reduction of known type.
   return false;
 }
 
-bool RecurrenceDescriptor::isFirstOrderRecurrence(PHINode *Phi, Loop *TheLoop,
-                                                  DominatorTree *DT) {
+bool RecurrenceDescriptor::isFirstOrderRecurrence(
+    PHINode *Phi, Loop *TheLoop,
+    DenseMap<Instruction *, Instruction *> &SinkAfter, DominatorTree *DT) {
 
   // Ensure the phi node is in the loop header and has two incoming values.
   if (Phi->getParent() != TheLoop->getHeader() ||
       Phi->getNumIncomingValues() != 2)
     return false;
 
   // Ensure the loop has a preheader and a single latch block. The loop
   // vectorizer will need the latch to set up the next iteration of the loop.
   auto *Preheader = TheLoop->getLoopPreheader();
   auto *Latch = TheLoop->getLoopLatch();
   if (!Preheader || !Latch)
     return false;
 
   // Ensure the phi node's incoming blocks are the loop preheader and latch.
   if (Phi->getBasicBlockIndex(Preheader) < 0 ||
       Phi->getBasicBlockIndex(Latch) < 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))
+  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.
+  if (Phi->hasOneUse()) {
+    auto *I = Phi->user_back();
+    if (I->isCast() && (I->getParent() == Phi->getParent()) && I->hasOneUse() &&
+        DT->dominates(Previous, I->user_back())) {
+      SinkAfter[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;
 }

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

public:
   ReductionList *getReductionVars() { return &Reductions; }
 
   /// Returns the induction variables found in the loop.
   InductionList *getInductionVars() { return &Inductions; }
 
   /// Return the first-order recurrences found in the loop.
   RecurrenceSet *getFirstOrderRecurrences() { return &FirstOrderRecurrences; }
 
+  /// Return the set of instructions to sink to handle first-order recurrences.
+  DenseMap<Instruction *, Instruction *> &getSinkAfter() { return SinkAfter; }
+
   /// Returns the widest induction type.
   Type *getWidestInductionType() { return WidestIndTy; }
 
   /// Returns True if V is an induction variable in this loop.
   bool isInductionVariable(const Value *V);
 
   /// Returns True if PN is a reduction variable in this loop.
   bool isReductionVariable(PHINode *PN) { return Reductions.count(PN); }
private:
   /// Holds the reduction variables.
   ReductionList Reductions;
   /// Holds all of the induction variables that we found in the loop.
   /// Notice that inductions don't need to start at zero and that induction
   /// variables can be pointers.
   InductionList Inductions;
   /// Holds the phi nodes that are first-order recurrences.
   RecurrenceSet FirstOrderRecurrences;
+  /// Holds instructions that need to sink past other instructions to handle
+  /// first-order recurrences.
+  DenseMap<Instruction *, Instruction *> SinkAfter;
   /// Holds the widest induction type encountered.
   Type *WidestIndTy;
 
   /// Allowed outside users. This holds the induction and reduction
   /// vars which can be accessed from outside the loop.
   SmallPtrSet<Value *, 4> AllowedExit;
 
   /// Can we assume the absence of NaNs.
for (Instruction &I : *BB) {
         InductionDescriptor ID;
         if (InductionDescriptor::isInductionPHI(Phi, TheLoop, PSE, ID)) {
           addInductionPhi(Phi, ID, AllowedExit);
           if (ID.hasUnsafeAlgebra() && !HasFunNoNaNAttr)
             Requirements->addUnsafeAlgebraInst(ID.getUnsafeAlgebraInst());
           continue;
         }
 
-        if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop, DT)) {
+        if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop,
+                                                         SinkAfter, DT)) {
           FirstOrderRecurrences.insert(Phi);
           continue;
         }
 
         // As a last resort, coerce the PHI to a AddRec expression
         // and re-try classifying it a an induction PHI.
         if (InductionDescriptor::isInductionPHI(Phi, TheLoop, PSE, ID, true)) {
           addInductionPhi(Phi, ID, AllowedExit);
void LoopVectorizationPlanner::executePlan(InnerLoopVectorizer &ILV) {
   // Notice: any optimization or new instruction that go
   // into the code below should also be implemented in
   // the cost-model.
   //
   //===------------------------------------------------===//
 
   // 2. Copy and widen instructions from the old loop into the new loop.
 
+  // Move instructions to handle first-order recurrences.
+  DenseMap<Instruction *, Instruction *> SinkAfter = Legal->getSinkAfter();
+  for (auto &Entry : SinkAfter) {
+    Entry.first->removeFromParent();
+    Entry.first->insertAfter(Entry.second);
+    DEBUG(dbgs() << "Sinking" << *Entry.first << " after" << *Entry.second
+                 << " to vectorize a 1st order recurrence.\n");
+  }
+
   // Collect instructions from the original loop that will become trivially dead
   // in the vectorized loop. We don't need to vectorize these instructions. For
   // example, original induction update instructions can become dead because we
   // separately emit induction "steps" when generating code for the new loop.
   // Similarly, we create a new latch condition when setting up the structure
   // of the new loop, so the old one can become dead.
   SmallPtrSet<Instruction *, 4> DeadInstructions;
   collectTriviallyDeadInstructions(DeadInstructions);

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

 ; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -dce -instcombine -S | FileCheck %s
 ; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=2 -dce -instcombine -S | FileCheck %s --check-prefix=UNROLL
 ; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=2 -S | FileCheck %s --check-prefix=UNROLL-NO-IC
 ; RUN: opt < %s -loop-vectorize -force-vector-width=1 -force-vector-interleave=2 -S | FileCheck %s --check-prefix=UNROLL-NO-VF
+; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S | FileCheck %s --check-prefix=SINK-AFTER
+; RUN: opt < %s -loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -S | FileCheck %s --check-prefix=NO-SINK-AFTER
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 
 ; void recurrence_1(int *a, int *b, int n) {
 ;   for(int i = 0; i < n; i++)
 ;     b[i] =  a[i] + a[i - 1]
 ; }
 ;
for.body:
   %tobool = fcmp une double %mul, 0.000000e+00
   %inc = zext i1 %tobool to i32
   %a.1 = add nsw i32 %a.010, %inc
   %inc1 = add nuw nsw i32 %i.011, 1
   %add.ptr = getelementptr inbounds double, double* %b.addr.012, i64 25
   %exitcond = icmp eq i32 %inc1, 10240
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 }
+
+; void sink_after(short *a, int n, int *b) {
+;   for(int i = 0; i < n; i++)
+;     b[i] = (a[i] * a[i + 1]);
+; }
+;
+; SINK-AFTER-LABEL: sink_after
+; Check that the sext sank after the load in the vector loop.
+; SINK-AFTER: vector.body
+; SINK-AFTER:   %vector.recur = phi <4 x i16> [ %vector.recur.init, %vector.ph ], [ %wide.load, %vector.body ]
+; SINK-AFTER:   %wide.load = load <4 x i16>
+; SINK-AFTER:   %[[VSHUF:.+]] = shufflevector <4 x i16> %vector.recur, <4 x i16> %wide.load, <4 x i32> <i32 3, i32 4, i32 5, i32 6>
+; SINK-AFTER:   %[[VCONV:.+]] = sext <4 x i16> %[[VSHUF]] to <4 x i32>
+; SINK-AFTER:   %[[VCONV3:.+]] = sext <4 x i16> %wide.load to <4 x i32>
+; SINK-AFTER:   mul nsw <4 x i32> %[[VCONV3]], %[[VCONV]]
+; Check also that the sext sank after the load in the scalar loop.
+; SINK-AFTER: for.body
+; SINK-AFTER:   %scalar.recur = phi i16 [ %scalar.recur.init, %scalar.ph ], [ %[[LOAD:.+]], %for.body ]
+; SINK-AFTER:   %[[LOAD]] = load i16, i16* %arrayidx2
+; SINK-AFTER:   %[[CONV:.+]] = sext i16 %scalar.recur to i32
+; SINK-AFTER:   %[[CONV3:.+]] = sext i16 %[[LOAD]] to i32
+; SINK-AFTER:   %mul = mul nsw i32 %[[CONV3]], %[[CONV]]
+;
+define void @sink_after(i16* %a, i32* %b, i64 %n) {
+entry:
+  %.pre = load i16, i16* %a
+  br label %for.body
+
+for.body:
+  %0 = phi i16 [ %.pre, %entry ], [ %1, %for.body ]
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %conv = sext i16 %0 to i32
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %arrayidx2 = getelementptr inbounds i16, i16* %a, i64 %indvars.iv.next
+  %1 = load i16, i16* %arrayidx2
+  %conv3 = sext i16 %1 to i32
+  %mul = mul nsw i32 %conv3, %conv
+  %arrayidx5 = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  store i32 %mul, i32* %arrayidx5
+  %exitcond = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}
+
+; void no_sink_after(short *a, int n, int *b) {
+;   for(int i = 0; i < n; i++)
+;     b[i] = ((a[i] + 2) * a[i + 1]);
+; }
+;
+; NO-SINK-AFTER-LABEL: no_sink_after
+; NO-SINK-AFTER-NOT:   vector.ph:
+; NO-SINK-AFTER:       }
+;
+define void @no_sink_after(i16* %a, i32* %b, i64 %n) {
+entry:
+  %.pre = load i16, i16* %a
+  br label %for.body
+
+for.body:
+  %0 = phi i16 [ %.pre, %entry ], [ %1, %for.body ]
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %conv = sext i16 %0 to i32
+  %add = add nsw i32 %conv, 2
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %arrayidx2 = getelementptr inbounds i16, i16* %a, i64 %indvars.iv.next
+  %1 = load i16, i16* %arrayidx2
+  %conv3 = sext i16 %1 to i32
+  %mul = mul nsw i32 %add, %conv3
+  %arrayidx5 = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  store i32 %mul, i32* %arrayidx5
+  %exitcond = icmp eq i64 %indvars.iv.next, %n
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}