[llvm-branch-commits] [llvm] [VPlan] Scalarize to first-lane-only directly on VPlan (PR #184267)

[Ramkumar Ramachandra via llvm-branch-commits]

================
@@ -6395,3 +6395,68 @@ void VPlanTransforms::makeMemOpWideningDecisions(
                              return ReplaceWith(VPI, Recipe);
                            });
 }
+
+void VPlanTransforms::makeScalarizationDecisions(
+    VPlan &Plan, VFRange &Range, VPRecipeBuilder &RecipeBuilder) {
+  if (LoopVectorizationPlanner::getDecisionAndClampRange(
+          [&](ElementCount VF) { return VF.isScalar(); }, Range))
+    return;
+
+  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+  VPBasicBlock *HeaderVPBB = LoopRegion->getEntryBasicBlock();
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           post_order<VPBlockShallowTraversalWrapper<VPBlockBase *>>(
+               HeaderVPBB))) {
+    for (VPRecipeBase &R :
+         make_early_inc_range(make_range(VPBB->rbegin(), VPBB->rend()))) {
+      auto *VPI = dyn_cast<VPInstruction>(&R);
+      if (!VPI)
+        continue;
+
+      auto *I = cast_or_null<Instruction>(VPI->getUnderlyingValue());
+      if (!I)
+        // Wouldn't be able to create a `VPReplicateRecipe` anyway.
+        continue;
+
+      bool CanTransformToFirstLaneOnly = [&]() {
+        if (VPI->mayHaveSideEffects())
+          return false;
+
+        if (is_contained({Instruction::SDiv, Instruction::UDiv,
+                          Instruction::SRem, Instruction::URem},
+                         VPI->getOpcode()) &&
+            VPI->getMask())
+          return false;
+
+        // Avoid rewriting IV increment as that interferes with
+        // `removeRedundantCanonicalIVs`.
+        if (VPI->getOpcode() == Instruction::Add &&
+            any_of(VPI->operands(),
+                   [&](auto *Op) { return isa<VPWidenInductionRecipe>(Op); }))
+          return false;
+
+        if (!all_of(VPI->users(), [&](auto *U) {
+              // TODO: This "ScalarCast" is bonkers...
+              if (VPI->isScalarCast() && isa<VPWidenGEPRecipe>(U))
+                return false;
+
+              return U->usesFirstLaneOnly(VPI);
+            }))
+          return false;
----------------
artagnon wrote:

vputils::onlyFirstLaneUsed?

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