[llvm] [VPlan] Explicitly unroll replicate-regions without live-outs by VF. (PR #170212)

[Florian Hahn via llvm-commits]

================
@@ -729,3 +733,150 @@ void VPlanTransforms::replicateByVF(VPlan &Plan, ElementCount VF) {
   for (auto *R : reverse(ToRemove))
     R->eraseFromParent();
 }
+
+/// Convert recipes in region blocks to operate on a single 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.
+static void convertRecipesInRegionBlocksToSingleScalar(
+    VPlan &Plan, Type *IdxTy, ElementCount VF,
+    ArrayRef<VPBlockBase *> RegionBlocks) {
+  for (VPBlockBase *VPB : RegionBlocks) {
+    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.
+        auto *DefR = Op->getDefiningRecipe();
+        if (VF.isScalar() || (DefR && DefR->getParent() == VPB) ||
+            vputils::isSingleScalar(Op))
+          continue;
+
+        // Extract the lane from values defined outside the region.
+        VPValue *Idx = Plan.getConstantInt(IdxTy, 0);
+        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();
+      }
+    }
+  }
+}
+
+/// Process recipes in a single lane's blocks, updating them for lane-specific
+/// operations.
----------------
fhahn wrote:

I think you were referring to `Old2NewBlocks`, right? we can just pass an arrayref of the new blocks, and simplify the loop using zip_equal. Hopefully it should be clearer now, thanks

https://github.com/llvm/llvm-project/pull/170212