Index: llvm/trunk/include/llvm/Transforms/Utils/LoopUtils.h =================================================================== --- llvm/trunk/include/llvm/Transforms/Utils/LoopUtils.h +++ llvm/trunk/include/llvm/Transforms/Utils/LoopUtils.h @@ -184,9 +184,14 @@ /// 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 &SinkAfter, + DominatorTree *DT); RecurrenceKind getRecurrenceKind() { return Kind; } Index: llvm/trunk/lib/Transforms/Utils/LoopUtils.cpp =================================================================== --- llvm/trunk/lib/Transforms/Utils/LoopUtils.cpp +++ llvm/trunk/lib/Transforms/Utils/LoopUtils.cpp @@ -528,8 +528,9 @@ return false; } -bool RecurrenceDescriptor::isFirstOrderRecurrence(PHINode *Phi, Loop *TheLoop, - DominatorTree *DT) { +bool RecurrenceDescriptor::isFirstOrderRecurrence( + PHINode *Phi, Loop *TheLoop, + DenseMap &SinkAfter, DominatorTree *DT) { // Ensure the phi node is in the loop header and has two incoming values. if (Phi->getParent() != TheLoop->getHeader() || @@ -551,12 +552,24 @@ // 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(Phi->getIncomingValueForBlock(Latch)); - if (!Previous || !TheLoop->contains(Previous) || isa(Previous)) + if (!Previous || !TheLoop->contains(Previous) || isa(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(U)) { if (!DT->dominates(Previous, I)) Index: llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp =================================================================== --- llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp +++ llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -1604,6 +1604,9 @@ /// 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 &getSinkAfter() { return SinkAfter; } + /// Returns the widest induction type. Type *getWidestInductionType() { return WidestIndTy; } @@ -1806,6 +1809,9 @@ 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 SinkAfter; /// Holds the widest induction type encountered. Type *WidestIndTy; @@ -5378,7 +5384,8 @@ continue; } - if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop, DT)) { + if (RecurrenceDescriptor::isFirstOrderRecurrence(Phi, TheLoop, + SinkAfter, DT)) { FirstOrderRecurrences.insert(Phi); continue; } @@ -7651,6 +7658,15 @@ // 2. Copy and widen instructions from the old loop into the new loop. + // Move instructions to handle first-order recurrences. + DenseMap 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 Index: llvm/trunk/test/Transforms/LoopVectorize/first-order-recurrence.ll =================================================================== --- llvm/trunk/test/Transforms/LoopVectorize/first-order-recurrence.ll +++ llvm/trunk/test/Transforms/LoopVectorize/first-order-recurrence.ll @@ -2,6 +2,8 @@ ; 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" @@ -447,3 +449,81 @@ %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> +; 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 +}