[VPlan] Explicitly unoll replicate-regions without live-outs by VF.

This patch adds a new replicateReplicateRegionsByVF transform to
unroll replicate=regions by VF, dissolving them. The transform creates
VF copies of the replicate-region's content, connects them and converts
recipes to single-scalar variants for the corresponding lanes.

The initial version skips regions with live-outs (VPPredInstPHIRecipe),
which will be added  in follow-up patches.

Depends on llvm#170053

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7306,6 +7306,10 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
                            BestVPlan);
   VPlanTransforms::runPass(VPlanTransforms::materializeBroadcasts, BestVPlan);
   VPlanTransforms::runPass(VPlanTransforms::replicateByVF, BestVPlan, BestVF);
+  VPlanTransforms::runPass(VPlanTransforms::unrollReplicateRegions, BestVPlan,
+                           BestVF);
+  VPlanTransforms::runPass(VPlanTransforms::mergeBlocksIntoPredecessors,
+                           BestVPlan);
   bool HasBranchWeights =
       hasBranchWeightMD(*OrigLoop->getLoopLatch()->getTerminator());
   if (HasBranchWeights) {

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3816,8 +3816,9 @@ class LLVM_ABI_FOR_TEST VPScalarIVStepsRecipe : public VPRecipeWithIRFlags,
         getOperand(0), getOperand(1), getOperand(2), InductionOpcode,
         hasFastMathFlags() ? getFastMathFlags() : FastMathFlags(),
         getDebugLoc());
-    if (getNumOperands() == 4)
-      NewR->addOperand(getOperand(3));
+    // Add lane/unroll-part operands, if present.
+    for (VPValue *Op : drop_begin(operands(), 3))
+      NewR->addOperand(Op);
     return NewR;
   }
 

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -481,7 +481,7 @@ static void addReplicateRegions(VPlan &Plan) {
 
 /// Remove redundant VPBasicBlocks by merging them into their predecessor if
 /// the predecessor has a single successor.
-static bool mergeBlocksIntoPredecessors(VPlan &Plan) {
+bool VPlanTransforms::mergeBlocksIntoPredecessors(VPlan &Plan) {
   SmallVector<VPBasicBlock *> WorkList;
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_deep(Plan.getEntry()))) {

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -178,6 +178,10 @@ struct VPlanTransforms {
   /// replicate regions, thereby dissolving the latter.
   static void replicateByVF(VPlan &Plan, ElementCount VF);
 
+  /// Replace replicate regions by explicitly replicating the regions' contents
+  /// \p VF times, each copy processing a single lane.
+  static void unrollReplicateRegions(VPlan &Plan, ElementCount VF);
+
   /// Optimize \p Plan based on \p BestVF and \p BestUF. This may restrict the
   /// resulting plan to \p BestVF and \p BestUF.
   static void optimizeForVFAndUF(VPlan &Plan, ElementCount BestVF,
@@ -189,6 +193,8 @@ struct VPlanTransforms {
   /// block merging.
   LLVM_ABI_FOR_TEST static void optimize(VPlan &Plan);
 
+  static bool mergeBlocksIntoPredecessors(VPlan &Plan);
+
   /// Wrap predicated VPReplicateRecipes with a mask operand in an if-then
   /// region block and remove the mask operand. Optimize the created regions by
   /// iteratively sinking scalar operands into the region, followed by merging

llvm/lib/Transforms/Vectorize/VPlanUnroll.cpp

@@ -24,6 +24,9 @@
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "vplan"
 
 using namespace llvm;
 using namespace llvm::VPlanPatternMatch;
@@ -121,6 +124,7 @@ class UnrollState {
       R->setOperand(OpIdx, getValueForPart(Op, Part));
   }
 };
+
 } // namespace
 
 void UnrollState::unrollReplicateRegionByUF(VPRegionBlock *VPR) {
@@ -634,3 +638,177 @@ void VPlanTransforms::replicateByVF(VPlan &Plan, ElementCount VF) {
   for (auto *R : reverse(ToRemove))
     R->eraseFromParent();
 }
+
+/// Process recipes in a single lane's blocks, updating them for lane-specific
+/// operations.
+static void processLane(VPlan &Plan, Type *IdxTy, unsigned Lane,
+                        ElementCount VF, ArrayRef<VPBlockBase *> RegionBlocks,
+                        DenseMap<VPBlockBase *, VPBlockBase *> &Old2NewBlocks) {
+  DenseMap<VPValue *, VPValue *> Old2NewVPValues;
+  for (VPBlockBase *OldVPB : RegionBlocks) {
+    auto *OldBB = cast<VPBasicBlock>(OldVPB);
+    auto *NewBB = cast<VPBasicBlock>(Old2NewBlocks.lookup(OldVPB));
+    for (const auto &[OldR, NewR] : zip(*OldBB, *NewBB)) {
+      for (const auto &[OldV, NewV] :
+           zip(OldR.definedValues(), NewR.definedValues()))
+        Old2NewVPValues[OldV] = NewV;
+    }
+
+    // Update lane operands and remap operands to use copies for current lane.
+    for (VPRecipeBase &NewR : make_early_inc_range(*NewBB)) {
+      if (auto *Steps = dyn_cast<VPScalarIVStepsRecipe>(&NewR))
+        Steps->setOperand(3, Plan.getConstantInt(IdxTy, Lane));
+      else if (match(&NewR, m_ExtractElement(m_VPValue(), m_ZeroInt())))
+        NewR.setOperand(1, Plan.getConstantInt(IdxTy, Lane));
+
+      // Remap operands to use lane-specific values.
+      for (const auto &[I, Op] : enumerate(NewR.operands())) {
+        // Use cloned value if operand was defined in the region.
+        if (auto *New = Old2NewVPValues.lookup(Op))
+          NewR.setOperand(I, New);
+      }
+    }
+  }
+}
+
+/// Process a single lane: clone blocks (or reuse original for lane 0), collect
+/// value mappings, and process recipes for lane-specific operations.
+static void processSingleLane(
+    VPlan &Plan, Type *IdxTy, unsigned Lane, ElementCount VF,
+    ArrayRef<VPBlockBase *> RegionBlocks, VPBlockBase *Entry,
+    VPBlockBase *Exiting,
+    SmallVectorImpl<std::pair<VPBlockBase *, VPBlockBase *>> &LaneClones) {
+  DenseMap<VPBlockBase *, VPBlockBase *> Old2NewBlocks;
+  if (Lane == 0) {
+    // Lane 0 uses the original blocks, and the recipes are adjusted:
+    // VPReplicateRecipes are converted to single-scalar ones, branch-on-mask is
+    // converted into BranchOnCond and extracts are created as needed.
+    for (VPBlockBase *VPB : RegionBlocks) {
+      Old2NewBlocks[VPB] = VPB;
+
+      for (VPRecipeBase &NewR :
+           make_early_inc_range(*cast<VPBasicBlock>(VPB))) {
+        VPBuilder Builder(&NewR);
+        for (const auto &[I, Op] : enumerate(NewR.operands())) {
+          // Skip operands that don't need extraction: scalar VF (no vectors),
+          // values defined in the same block (already scalar), or values that
+          // are already single scalars.
+          if (VF.isScalar() ||
+              (Op->getDefiningRecipe() &&
+               Op->getDefiningRecipe()->getParent() == VPB) ||
+              vputils::isSingleScalar(Op))
+            continue;
+
+          // Extract the lane from values defined outside the region.
+          VPValue *Idx = Plan.getConstantInt(IdxTy, Lane);
+          VPValue *Extract = Builder.createNaryOp(
+              Instruction::ExtractElement, {Op, Idx}, NewR.getDebugLoc());
+          NewR.setOperand(I, Extract);
+        }
+
+        if (auto *RepR = dyn_cast<VPReplicateRecipe>(&NewR)) {
+          auto *New = new VPReplicateRecipe(
+              RepR->getUnderlyingInstr(), RepR->operands(),
+              /* IsSingleScalar=*/true, /*Mask=*/nullptr, *RepR, *RepR,
+              RepR->getDebugLoc());
+          New->insertBefore(RepR);
+          RepR->replaceAllUsesWith(New);
+          RepR->eraseFromParent();
+        } else if (auto *BranchOnMask = dyn_cast<VPBranchOnMaskRecipe>(&NewR)) {
+          Builder.createNaryOp(VPInstruction::BranchOnCond,
+                               {BranchOnMask->getOperand(0)},
+                               BranchOnMask->getDebugLoc());
+          BranchOnMask->eraseFromParent();
+        } else if (auto *Steps = dyn_cast<VPScalarIVStepsRecipe>(&NewR)) {
+          // Add lane operand (4th operand) for VPScalarIVStepsRecipe if not
+          // already present.
+          unsigned NumOps = Steps->getNumOperands();
+          if (NumOps == 4) {
+            // Has UnrollPart at position 3, need to insert Lane before it.
+            VPValue *UnrollPart = Steps->getOperand(3);
+            Steps->setOperand(3, Plan.getConstantInt(IdxTy, Lane));
+            Steps->addOperand(UnrollPart);
+          } else if (NumOps == 3) {
+            // Just BaseIV, Step, VF - add Lane.
+            Steps->addOperand(Plan.getConstantInt(IdxTy, Lane));
+            Steps->addOperand(Plan.getConstantInt(IdxTy, 0));
+          }
+        }
+      }
+    }
+  } else {
+    // Clone blocks and connect them according to original structure.
+    for (VPBlockBase *OrigBlock : RegionBlocks) {
+      VPBlockBase *ClonedBlock = OrigBlock->clone();
+      Old2NewBlocks[OrigBlock] = ClonedBlock;
+      ClonedBlock->setParent(Entry->getParent());
+    }
+    for (VPBlockBase *OrigBlock : RegionBlocks) {
+      if (OrigBlock == Exiting)
+        continue;
+      for (VPBlockBase *OrigSucc : OrigBlock->successors())
+        VPBlockUtils::connectBlocks(Old2NewBlocks[OrigBlock],
+                                    Old2NewBlocks[OrigSucc]);
+    }
+  }
+
+  processLane(Plan, IdxTy, Lane, VF, RegionBlocks, Old2NewBlocks);
+  LaneClones.push_back({Old2NewBlocks[Entry], Old2NewBlocks[Exiting]});
+}
+
+void VPlanTransforms::unrollReplicateRegions(VPlan &Plan, ElementCount VF) {
+  // Collect all replicate regions in the plan before modifying the CFG.
+  SmallVector<VPRegionBlock *> ReplicateRegions;
+  for (VPBlockBase *Block :
+       vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry())) {
+    if (auto *Region = dyn_cast<VPRegionBlock>(Block)) {
+      if (Region->isReplicator())
+        ReplicateRegions.push_back(Region);
+    }
+  }
+
+  Type *IdxTy = IntegerType::get(Plan.getContext(), 32);
+
+  for (VPRegionBlock *Region : ReplicateRegions) {
+    assert(!VF.isScalable() && "cannot replicate across scalable VFs");
+
+    VPBlockBase *Entry = Region->getEntry();
+    VPBlockBase *Exiting = Region->getExiting();
+
+    // Skip regions with live-outs as packing scalar results back into vectors
+    // is not yet implemented.
+    if (any_of(*cast<VPBasicBlock>(Exiting), IsaPred<VPPredInstPHIRecipe>))
+      continue;
+
+    // Get region context before dissolving.
+    VPBlockBase *Pred = Region->getSinglePredecessor();
+    assert(Pred && "Replicate region must have a single predecessor");
+    SmallVector<VPBlockBase *> Successors(Region->successors());
+
+    // Disconnect and dissolve the region.
+    VPBlockUtils::disconnectBlocks(Pred, Region);
+    for (VPBlockBase *Succ : Successors)
+      VPBlockUtils::disconnectBlocks(Region, Succ);
+
+    SmallVector<VPBlockBase *> RegionBlocks(vp_depth_first_shallow(Entry));
+    VPRegionBlock *ParentRegion = Region->getParent();
+    for (VPBlockBase *Block : RegionBlocks)
+      Block->setParent(ParentRegion);
+    VPBlockUtils::connectBlocks(Pred, Entry);
+
+    // Process each lane: clone blocks, collect value mappings, and process
+    // recipes for lane-specific operations.
+    SmallVector<std::pair<VPBlockBase *, VPBlockBase *>> LaneClones;
+    for (unsigned Lane = 0; Lane < VF.getKnownMinValue(); ++Lane) {
+      processSingleLane(Plan, IdxTy, Lane, VF, RegionBlocks, Entry, Exiting,
+                        LaneClones);
+    }
+
+    // Connect lanes sequentially and connect last lane to successors.
+    for (unsigned Lane = 1; Lane < VF.getKnownMinValue(); ++Lane)
+      VPBlockUtils::connectBlocks(LaneClones[Lane - 1].second,
+                                  LaneClones[Lane].first);
+    for (VPBlockBase *Succ : Successors)
+      VPBlockUtils::connectBlocks(LaneClones.back().second, Succ);
+  }
+}

llvm/test/Transforms/LoopVectorize/AArch64/blend-costs.ll

@@ -199,8 +199,7 @@ define void @test_blend_feeding_replicated_store_2(ptr noalias %src, ptr %dst, i
 ; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <16 x i1> [[TMP7]], i32 0
 ; CHECK-NEXT:    br i1 [[TMP8]], label %[[PRED_STORE_IF:.*]], label %[[PRED_STORE_CONTINUE:.*]]
 ; CHECK:       [[PRED_STORE_IF]]:
-; CHECK-NEXT:    [[TMP72:%.*]] = add i32 [[IV]], 0
-; CHECK-NEXT:    [[TMP9:%.*]] = getelementptr inbounds i8, ptr [[DST]], i32 [[TMP72]]
+; CHECK-NEXT:    [[TMP9:%.*]] = getelementptr inbounds i8, ptr [[DST]], i32 [[IV]]
 ; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <16 x i8> [[PREDPHI]], i32 0
 ; CHECK-NEXT:    store i8 [[TMP10]], ptr [[TMP9]], align 1
 ; CHECK-NEXT:    br label %[[PRED_STORE_CONTINUE]]