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

[Luke Lau via llvm-commits]

================
@@ -6541,3 +6541,60 @@ void VPlanTransforms::makeMemOpWideningDecisions(
     ReplaceWith(Recipe);
   }
 }
+
+void VPlanTransforms::makeScalarizationDecisions(VPlan &Plan, VFRange &Range) {
+  if (LoopVectorizationPlanner::getDecisionAndClampRange(
+          [&](ElementCount VF) { return VF.isScalar(); }, Range))
+    return;
+
+  // Extend lifetime per `llvm::PostOrderTraversal` documentation:
+  auto PO =
+      vp_post_order_shallow(Plan.getVectorLoopRegion()->getEntryBasicBlock());
+
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(PO)) {
+    for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
+      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;
+
+        // We want to drop the mask operand, doing that for integer division
+        // isn't safe.
+        if (is_contained({Instruction::SDiv, Instruction::UDiv,
----------------
lukel97 wrote:

Ah right I see, we're also checking if there's a mask. That makes sense

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