[VPlan] Sink non-uniform recieps for scalar plans.

[fhahn]

Summary

In scalar plans, replicate recipes will only generate a single value per
UF, independent of whether they are uniform or not. So don't consider
uniformity for plans with scalar VFs only.

This allows us to handle a few additional cases in VPlan sinking instead
of non-VPlan sinkScalarOperands.

Depends on D133762.

[Ayal]

LGTM.

This is essentially saying that for VF=1 every recipe is uniform (across VF lanes), or none are.
Perhaps this should be done more 'uniformly', as in CM.isUniformAfterVectorization(I, VF) which returns true if (VF.isScalar()) regardless of I.

llvm/lib/Transforms/Vectorize/VPlan.h

2674

nit: might also be useful when printing VF range of VPlan, currently set in VPlan's name.

[fhahn]

Rebased patch on top of current main.

In D134218#3802359, @Ayal wrote:

LGTM.

This is essentially saying that for VF=1 every recipe is uniform (across VF lanes), or none are.
Perhaps this should be done more 'uniformly', as in CM.isUniformAfterVectorization(I, VF) which returns true if (VF.isScalar()) regardless of I.

Thanks! That might be an option in the future.

llvm/lib/Transforms/Vectorize/VPlan.h

public:
   /// Mark the plan to indicate that using Value2VPValue is not safe any
   /// longer, because it may be stale.
   void disableValue2VPValue() { Value2VPValueEnabled = false; }
 
   void addVF(ElementCount VF) { VFs.insert(VF); }
 
   bool hasVF(ElementCount VF) { return VFs.count(VF); }
 
+  bool hasScalarVFOnly() const { return VFs.size() == 1 && VFs[0].isScalar(); }
+
   const std::string &getName() const { return Name; }
 
   void setName(const Twine &newName) { Name = newName.str(); }
 
   /// Get the existing or add a new external definition for \p V.
   VPValue *getOrAddExternalDef(Value *V) {
     auto I = VPExternalDefs.insert({V, nullptr});
     if (I.second)

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (auto &Recipe : *VPBB) {
       if (!RepR || !RepR->isPredicated())
         continue;
       for (VPValue *Op : RepR->operands())
         if (auto *Def = Op->getDefiningRecipe())
           WorkList.insert(std::make_pair(RepR->getParent(), Def));
     }
   }
 
+  bool ScalarVFOnly = Plan.hasScalarVFOnly();
   // Try to sink each replicate or scalar IV steps recipe in the worklist.
   while (!WorkList.empty()) {
     VPBasicBlock *SinkTo;
     VPRecipeBase *SinkCandidate;
     std::tie(SinkTo, SinkCandidate) = WorkList.pop_back_val();
     if (SinkCandidate->getParent() == SinkTo ||
         SinkCandidate->mayHaveSideEffects() ||
         SinkCandidate->mayReadOrWriteMemory())
       continue;
     if (auto *RepR = dyn_cast<VPReplicateRecipe>(SinkCandidate)) {
-      if (RepR->isUniform())
+      if (!ScalarVFOnly && RepR->isUniform())
         continue;
     } else if (!isa<VPScalarIVStepsRecipe>(SinkCandidate))
       continue;
 
     bool NeedsDuplicating = false;
     // All recipe users of the sink candidate must be in the same block SinkTo
     // or all users outside of SinkTo must be uniform-after-vectorization (
     // i.e., only first lane is used) . In the latter case, we need to duplicate
auto CanSinkWithUser = [SinkTo, &NeedsDuplicating,
           UI->onlyFirstLaneUsed(SinkCandidate->getVPSingleValue());
       // We only know how to duplicate VPRecipeRecipes for now.
       return NeedsDuplicating && isa<VPReplicateRecipe>(SinkCandidate);
     };
     if (!all_of(SinkCandidate->getVPSingleValue()->users(), CanSinkWithUser))
       continue;
 
     if (NeedsDuplicating) {
+      if (ScalarVFOnly)
+        continue;
       Instruction *I = cast<Instruction>(
           cast<VPReplicateRecipe>(SinkCandidate)->getUnderlyingValue());
       auto *Clone =
           new VPReplicateRecipe(I, SinkCandidate->operands(), true, false);
       // TODO: add ".cloned" suffix to name of Clone's VPValue.
 
       Clone->insertBefore(SinkCandidate);
       SmallVector<VPUser *, 4> Users(

llvm/test/Transforms/LoopVectorize/interleave-and-scalarize-only.ll

 ; DBG-NEXT:   vector.body:
 ; DBG-NEXT:     EMIT vp<[[CAN_IV:%.+]]> = CANONICAL-INDUCTION
 ; DBG-NEXT:     vp<[[DERIVED_IV:%.+]]> = DERIVED-IV ir<false> + vp<[[CAN_IV]]> * ir<true>
 ; DBG-NEXT:     vp<[[STEPS1:%.+]]>    = SCALAR-STEPS vp<[[DERIVED_IV]]>, ir<true>
 ; DBG-NEXT:     vp<[[STEPS2:%.+]]>    = SCALAR-STEPS vp<[[CAN_IV]]>, ir<1>
 ; DBG-NEXT:   Successor(s): cond.false
 ; DBG-EMPTY:
 ; DBG-NEXT:   cond.false:
-; DBG-NEXT:     CLONE ir<%gep.src> = getelementptr ir<%src>, vp<[[STEPS2]]>
-; DBG-NEXT:     CLONE ir<%gep.dst> = getelementptr ir<%dst>, vp<[[STEPS2]]>
 ; DBG-NEXT:   Successor(s): cond.false.0
 ; DBG-EMPTY:
 ; DBG-NEXT:   cond.false.0:
 ; DBG-NEXT:   Successor(s): pred.store
 ; DBG-EMPTY:
 ; DBG-NEXT:   <xVFxUF> pred.store: {
 ; DBG-NEXT:     pred.store.entry:
 ; DBG-NEXT:       BRANCH-ON-MASK vp<[[STEPS1]]>
 ; DBG-NEXT:     Successor(s): pred.store.if, pred.store.continue
 ; DBG-EMPTY:
 ; DBG-NEXT:     pred.store.if:
+; DBG-NEXT:       CLONE ir<%gep.src> = getelementptr ir<%src>, vp<[[STEPS2]]>
 ; DBG-NEXT:       CLONE ir<%l> = load ir<%gep.src>
+; DBG-NEXT:       CLONE ir<%gep.dst> = getelementptr ir<%dst>, vp<[[STEPS2]]>
 ; DBG-NEXT:       CLONE store ir<%l>, ir<%gep.dst>
 ; DBG-NEXT:     Successor(s): pred.store.continue
 ; DBG-EMPTY:
 ; DBG-NEXT:     pred.store.continue:
 ; DBG-NEXT:       PHI-PREDICATED-INSTRUCTION vp<{{.+}}> = ir<%l>
 ; DBG-NEXT:     No successors
 ; DBG-NEXT:   }
 ; DBG-NEXT:   Successor(s): cond.false.1

 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %pred.store.continue5 ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[INDEX]] to i1
 ; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = sub i1 false, [[TMP0]]
 ; CHECK-NEXT:    [[INDUCTION:%.*]] = add i1 [[OFFSET_IDX]], false
 ; CHECK-NEXT:    [[INDUCTION3:%.*]] = add i1 [[OFFSET_IDX]], true
 ; CHECK-NEXT:    br i1 [[INDUCTION]], label %pred.store.if, label %pred.store.continue
 ; CHECK:       pred.store.if:
 ; CHECK-NEXT:    [[INDUCTION4:%.*]] = add i64 [[INDEX]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr %dst, i64 [[INDUCTION4]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, ptr %src, i64 [[INDUCTION4]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = load i32, ptr [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr %dst, i64 [[INDUCTION4]]
 ; CHECK-NEXT:    store i32 [[TMP4]], ptr [[TMP1]], align 4
 ; CHECK-NEXT:    br label %pred.store.continue
 ; CHECK:       pred.store.continue:
 ; CHECK-NEXT:    [[TMP5:%.*]] = phi i32 [ poison, %vector.body ], [ [[TMP4]], %pred.store.if ]
 ; CHECK-NEXT:    br i1 [[INDUCTION3]], label %pred.store.if4, label %pred.store.continue5
 ; CHECK:       pred.store.if4:
 ; CHECK-NEXT:    [[INDUCTION5:%.*]] = add i64 [[INDEX]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, ptr %dst, i64 [[INDUCTION5]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, ptr %src, i64 [[INDUCTION5]]
 ; CHECK-NEXT:    [[TMP7:%.*]] = load i32, ptr [[TMP6]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, ptr %dst, i64 [[INDUCTION5]]
 ; CHECK-NEXT:    store i32 [[TMP7]], ptr [[TMP2]], align 4
 ; CHECK-NEXT:    br label %pred.store.continue5
 ; CHECK:       pred.store.continue5:
 ; CHECK-NEXT:    [[TMP8:%.*]] = phi i32 [ poison, %pred.store.continue ], [ [[TMP7]], %pred.store.if4 ]
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 2
 ; CHECK-NEXT:    [[TMP9:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1000
 ; CHECK-NEXT:    br i1 [[TMP9]], label %middle.block, label %vector.body
 ; CHECK:       middle.block:

llvm/test/Transforms/LoopVectorize/vplan-sink-scalars-and-merge-vf1.ll

 ; CHECK-EMPTY:
 ; CHECK-NEXT: <x1> vector loop: {
 ; CHECK-NEXT: vector.body:
 ; CHECK-NEXT:   EMIT vp<[[CAN_IV:%.+]]> = CANONICAL-INDUCTION
 ; CHECK-NEXT:   vp<[[STEPS:%.+]]> = SCALAR-STEPS vp<[[CAN_IV]]>, ir<1>
 ; CHECK-NEXT:   CLONE ir<%tmp2> = getelementptr ir<%ptr>, vp<[[STEPS]]>
 ; CHECK-NEXT:   CLONE ir<%tmp3> = load ir<%tmp2>
 ; CHECK-NEXT:   CLONE store ir<0>, ir<%tmp2>
-; CHECK-NEXT:   CLONE ir<%tmp4> = zext ir<%tmp3>
-; CHECK-NEXT:   CLONE ir<%tmp5> = trunc ir<%tmp4>
 ; CHECK-NEXT: Successor(s): if.then
 
 ; CHECK:      if.then:
 ; CHECK-NEXT: Successor(s): pred.store
 
 ; CHECK:      <xVFxUF> pred.store: {
 ; CHECK-NEXT:   pred.store.entry:
 ; CHECK-NEXT:     BRANCH-ON-MASK ir<%c>
 ; CHECK-NEXT:   Successor(s): pred.store.if, pred.store.continue
 
 ; CHECK:      pred.store.if:
+; CHECK-NEXT:   CLONE ir<%tmp4> = zext ir<%tmp3>
+; CHECK-NEXT:   CLONE ir<%tmp5> = trunc ir<%tmp4>
 ; CHECK-NEXT:   CLONE store ir<%tmp5>, ir<%tmp2>
 ; CHECK-NEXT:   Successor(s): pred.store.continue
 
 ; CHECK:      pred.store.continue:
 ; CHECK-NEXT:   No successors
 ; CHECK-NEXT: }
 
 ; CHECK:      if.then.0: