!fixup address comments, thnaks

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7267,8 +7267,8 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
   VPlanTransforms::runPass(VPlanTransforms::unrollByUF, BestVPlan, BestUF,
                            OrigLoop->getHeader()->getContext());
   VPlanTransforms::runPass(VPlanTransforms::materializeBuildVectors, BestVPlan);
-  VPlanTransforms::runPass(VPlanTransforms::replicateByVF, BestVPlan, BestVF);
   VPlanTransforms::runPass(VPlanTransforms::materializeBroadcasts, BestVPlan);
+  VPlanTransforms::runPass(VPlanTransforms::replicateByVF, BestVPlan, BestVF);
   bool HasBranchWeights =
       hasBranchWeightMD(*OrigLoop->getLoopLatch()->getTerminator());
   if (HasBranchWeights) {

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -372,8 +372,9 @@ Value *VPTransformState::get(const VPValue *Def, bool NeedsScalar) {
     set(Def, VectorValue);
   } else {
     assert(!VF.isScalable() && "VF is assumed to be non scalable.");
-    assert(isa<VPInstruction>(Def) && "Explicit BuildVector recipes must "
-                                      "handle packing for non-VPInstructions.");
+    assert(isa<VPInstruction>(Def) &&
+           "Explicit BuildVector recipes must have"
+           "handled packing for non-VPInstructions.");
     // Initialize packing with insertelements to start from poison.
     VectorValue = PoisonValue::get(toVectorizedTy(LastInst->getType(), VF));
     for (unsigned Lane = 0; Lane < VF.getFixedValue(); ++Lane)

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -3289,38 +3289,37 @@ void VPlanTransforms::materializeBuildVectors(VPlan &Plan) {
   auto VPBBsInsideLoopRegion = VPBlockUtils::blocksOnly<VPBasicBlock>(
       vp_depth_first_shallow(LoopRegion->getEntry()));
   // Materialize Build(Struct)Vector for all replicating VPReplicateRecipes,
-  // excluding ones in replicate regions. Those are not unrolled explicitly yet.
+  // excluding ones in replicate regions. Those are not materialized explicitly
+  // yet. Those vector users are still handled in VPReplicateRegion::execute(),
+  // via shouldPack().
+  // TODO: materialize build vectors for replicating recipes in replicating
+  // regions.
+  // TODO: materialize build vectors for VPInstructions.
   for (VPBasicBlock *VPBB :
        concat<VPBasicBlock *>(VPBBsOutsideLoopRegion, VPBBsInsideLoopRegion)) {
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
       auto *RepR = dyn_cast<VPReplicateRecipe>(&R);
-      if (!RepR || RepR->isSingleScalar())
-        continue;
-      VPInstruction *BuildVector = nullptr;
-      for (VPUser *U : to_vector(RepR->users())) {
+      auto UsesVectorOrInsideReplicateRegions = [RepR, LoopRegion](VPUser *U) {
         VPRegionBlock *ParentRegion =
             cast<VPRecipeBase>(U)->getParent()->getParent();
-        if (U->usesScalars(RepR) && ParentRegion == LoopRegion)
-          continue;
-
-        if (!BuildVector) {
-          Type *ScalarTy = TypeInfo.inferScalarType(RepR);
-          unsigned Opc = ScalarTy->isStructTy()
-                             ? VPInstruction::BuildStructVector
-                             : VPInstruction::BuildVector;
-          BuildVector = new VPInstruction(Opc, {RepR});
-          BuildVector->insertAfter(RepR);
-        }
+        return !U->usesScalars(RepR) || ParentRegion != LoopRegion;
+      };
+      if (!RepR || RepR->isSingleScalar() ||
+          none_of(RepR->users(), UsesVectorOrInsideReplicateRegions))
+        continue;
 
-        // Only update a single operand per users, as the same user is added
-        // multiple times, once per use.
-        // TODO: Introduce de-duplicating iterator over users.
-        for (unsigned Idx = 0; Idx != U->getNumOperands(); ++Idx)
-          if (U->getOperand(Idx) == RepR) {
-            U->setOperand(Idx, BuildVector);
-            break;
-          }
-      }
+      Type *ScalarTy = TypeInfo.inferScalarType(RepR);
+      unsigned Opcode = ScalarTy->isStructTy()
+                            ? VPInstruction::BuildStructVector
+                            : VPInstruction::BuildVector;
+      auto *BuildVector = new VPInstruction(Opcode, {RepR});
+      BuildVector->insertAfter(RepR);
+
+      RepR->replaceUsesWithIf(
+          BuildVector, [BuildVector, &UsesVectorOrInsideReplicateRegions](
+                           VPUser &U, unsigned) {
+            return &U != BuildVector && UsesVectorOrInsideReplicateRegions(&U);
+          });
     }
   }
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -276,8 +276,8 @@ struct VPlanTransforms {
   static void materializeBackedgeTakenCount(VPlan &Plan,
                                             VPBasicBlock *VectorPH);
 
-  /// Add explicit Build[Struct]Vector recipes that combine scalar values
-  /// produced by VPReplicateRecipes to a single vector.
+  /// Add explicit Build[Struct]Vector recipes that combine multiple scalar
+  /// values into single vectors.
   static void materializeBuildVectors(VPlan &Plan);
 
   /// Try to convert a plan with interleave groups with VF elements to a plan

llvm/lib/Transforms/Vectorize/VPlanUnroll.cpp

@@ -465,11 +465,12 @@ void VPlanTransforms::unrollByUF(VPlan &Plan, unsigned UF, LLVMContext &Ctx) {
   VPlanTransforms::removeDeadRecipes(Plan);
 }
 
-/// Create a single-scalar clone of \p RepR for lane \p Lane.
+/// Create a single-scalar clone of \p RepR for lane \p Lane. Use \p
+/// Def2LaneDefs to look up scalar definitions for operands of \RepR.
 static VPReplicateRecipe *
 cloneForLane(VPlan &Plan, VPBuilder &Builder, Type *IdxTy,
              VPReplicateRecipe *RepR, VPLane Lane,
-             DenseMap<VPValue *, SmallVector<VPValue *>> &Value2Lanes) {
+             const DenseMap<VPValue *, SmallVector<VPValue *>> &Def2LaneDefs) {
   // Collect the operands at Lane, creating extracts as needed.
   SmallVector<VPValue *> NewOps;
   for (VPValue *Op : RepR->operands()) {
@@ -482,8 +483,11 @@ cloneForLane(VPlan &Plan, VPBuilder &Builder, Type *IdxTy,
           Builder.createNaryOp(VPInstruction::ExtractLastElement, {Op}));
       continue;
     }
-    if (Value2Lanes.contains(Op)) {
-      NewOps.push_back(Value2Lanes[Op][Lane.getKnownLane()]);
+    // If Op is a definition that has been unrolled, directly use the clone for
+    // the corresponding lane.
+    auto LaneDefs = Def2LaneDefs.find(Op);
+    if (LaneDefs != Def2LaneDefs.end()) {
+      NewOps.push_back(LaneDefs->second[Lane.getKnownLane()]);
       continue;
     }
 
@@ -519,7 +523,7 @@ void VPlanTransforms::replicateByVF(VPlan &Plan, ElementCount VF) {
       vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()));
   auto VPBBsToUnroll =
       concat<VPBasicBlock *>(VPBBsOutsideLoopRegion, VPBBsInsideLoopRegion);
-  DenseMap<VPValue *, SmallVector<VPValue *>> Value2Lanes;
+  DenseMap<VPValue *, SmallVector<VPValue *>> Def2LaneDefs;
   SmallVector<VPRecipeBase *> ToRemove;
   for (VPBasicBlock *VPBB : VPBBsToUnroll) {
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
@@ -533,11 +537,11 @@ void VPlanTransforms::replicateByVF(VPlan &Plan, ElementCount VF) {
             vputils::isSingleScalar(RepR->getOperand(1))) {
           // Stores to invariant addresses need to store the last lane only.
           cloneForLane(Plan, Builder, IdxTy, RepR, VPLane::getLastLaneForVF(VF),
-                       Value2Lanes);
+                       Def2LaneDefs);
         } else {
           // Create single-scalar version of RepR for all lanes.
           for (unsigned I = 0; I != VF.getKnownMinValue(); ++I)
-            cloneForLane(Plan, Builder, IdxTy, RepR, VPLane(I), Value2Lanes);
+            cloneForLane(Plan, Builder, IdxTy, RepR, VPLane(I), Def2LaneDefs);
         }
         RepR->eraseFromParent();
         continue;
@@ -546,22 +550,24 @@ void VPlanTransforms::replicateByVF(VPlan &Plan, ElementCount VF) {
       SmallVector<VPValue *> LaneDefs;
       for (unsigned I = 0; I != VF.getKnownMinValue(); ++I)
         LaneDefs.push_back(
-            cloneForLane(Plan, Builder, IdxTy, RepR, VPLane(I), Value2Lanes));
+            cloneForLane(Plan, Builder, IdxTy, RepR, VPLane(I), Def2LaneDefs));
 
-      Value2Lanes[RepR] = LaneDefs;
+      Def2LaneDefs[RepR] = LaneDefs;
       /// Users that only demand the first lane can use the definition for lane
       /// 0.
       RepR->replaceUsesWithIf(LaneDefs[0], [RepR](VPUser &U, unsigned) {
         return U.onlyFirstLaneUsed(RepR);
       });
 
+      // Update each build vector user that currently has RepR as its only
+      // operand, to have all LaneDefs as its operands.
       for (VPUser *U : to_vector(RepR->users())) {
         auto *VPI = dyn_cast<VPInstruction>(U);
         if (!VPI || (VPI->getOpcode() != VPInstruction::BuildVector &&
                      VPI->getOpcode() != VPInstruction::BuildStructVector))
           continue;
         assert(VPI->getNumOperands() == 1 &&
-               "Build(Struct)Vector must have a single operand");
+               "Build(Struct)Vector must have a single operand before replicating by VF"");
         VPI->setOperand(0, LaneDefs[0]);
         for (VPValue *Def : drop_begin(LaneDefs))
           VPI->addOperand(Def);