[VPlan] Handle VPReplicateRecipe stores

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h

@@ -364,6 +364,12 @@ m_Broadcast(const Op0_t &Op0) {
   return m_VPInstruction<VPInstruction::Broadcast>(Op0);
 }
 
+template <typename Op0_t>
+inline UnaryVPInstruction_match<Op0_t, VPInstruction::ExtractLastElement>
+m_ExtractLastElement(const Op0_t &Op0) {
+  return m_VPInstruction<VPInstruction::ExtractLastElement>(Op0);
+}
+
 template <typename Op0_t, typename Op1_t>
 inline BinaryVPInstruction_match<Op0_t, Op1_t, VPInstruction::ActiveLaneMask>
 m_ActiveLaneMask(const Op0_t &Op0, const Op1_t &Op1) {

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2654,14 +2654,6 @@ void VPReplicateRecipe::execute(VPTransformState &State) {
     return;
   }
 
-  // A store of a loop varying value to a uniform address only needs the last
-  // copy of the store.
-  if (isa<StoreInst>(UI) && vputils::isSingleScalar(getOperand(1))) {
-    auto Lane = VPLane::getLastLaneForVF(State.VF);
-    scalarizeInstruction(UI, this, VPLane(Lane), State);
-    return;
-  }
-
   // Generate scalar instances for all VF lanes.
   const unsigned EndLane = State.VF.getFixedValue();
   for (unsigned Lane = 0; Lane < EndLane; ++Lane)

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -683,7 +683,8 @@ static void legalizeAndOptimizeInductions(VPlan &Plan) {
     // scalars.
     auto *WideIV = cast<VPWidenIntOrFpInductionRecipe>(&Phi);
     if (HasOnlyVectorVFs && none_of(WideIV->users(), [WideIV](VPUser *U) {
-          return U->usesScalars(WideIV);
+          return U->usesScalars(WideIV) ||
+                 match(U, m_ExtractLastElement(m_VPValue()));
         }))
       continue;
 
@@ -694,12 +695,19 @@ static void legalizeAndOptimizeInductions(VPlan &Plan) {
         WideIV->getTruncInst(), WideIV->getStartValue(), WideIV->getStepValue(),
         WideIV->getDebugLoc(), Builder);
 
+    bool HasWideOps = any_of(WideIV->users(), [WideIV](VPUser *U) {
+      return !U->usesScalars(WideIV) &&
+             !match(U, m_ExtractLastElement(m_VPValue()));
+    });
+
     // Update scalar users of IV to use Step instead.
     if (!HasOnlyVectorVFs)
       WideIV->replaceAllUsesWith(Steps);
     else
-      WideIV->replaceUsesWithIf(Steps, [WideIV](VPUser &U, unsigned) {
-        return U.usesScalars(WideIV);
+      WideIV->replaceUsesWithIf(Steps, [WideIV, HasWideOps](VPUser &U,
+                                                            unsigned) {
+        return U.usesScalars(WideIV) ||
+               (!HasWideOps && match(&U, m_ExtractLastElement(m_VPValue())));
       });
   }
 }
@@ -1209,6 +1217,20 @@ static void narrowToSingleScalarRecipes(VPlan &Plan) {
         continue;
 
       auto *RepOrWidenR = cast<VPSingleDefRecipe>(&R);
+
+      if (RepR && isa<StoreInst>(RepR->getUnderlyingInstr()) &&
+          vputils::isSingleScalar(RepR->getOperand(1))) {
+        auto *Clone = new VPReplicateRecipe(
+            RepOrWidenR->getUnderlyingInstr(), RepOrWidenR->operands(),
+            true /*IsSingleScalar*/, nullptr /*Mask*/, *RepR /*Metadata*/);
+        Clone->insertBefore(RepOrWidenR);
+        auto *Ext = new VPInstruction(VPInstruction::ExtractLastElement,
+                                      {Clone->getOperand(0)});
+        Ext->insertBefore(Clone);
+        Clone->setOperand(0, Ext);
+        RepR->eraseFromParent();
+        continue;
+      }
       // Skip recipes that aren't single scalars or don't have only their
       // scalar results used. In the latter case, we would introduce extra
       // broadcasts.