[llvm] [LV] Convert scatter w/uniform addr and mask being header mask to scalar store. (PR #172799)

[Elvis Wang via llvm-commits]

================
@@ -1785,41 +1785,85 @@ static void reassociateHeaderMask(VPlan &Plan) {
   }
 }
 
-static void narrowToSingleScalarRecipes(VPlan &Plan) {
+void VPlanTransforms::narrowScatters(VPlan &Plan, VPCostContext &Ctx,
+                                     VFRange &Range,
+                                     const bool &FoldTailWithEVL) {
   if (Plan.hasScalarVFOnly())
     return;
 
-  // Try to narrow wide and replicating recipes to single scalar recipes,
-  // based on VPlan analysis. Only process blocks in the loop region for now,
-  // without traversing into nested regions, as recipes in replicate regions
-  // cannot be converted yet.
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
            vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()))) {
     for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
-      if (!isa<VPWidenRecipe, VPWidenGEPRecipe, VPReplicateRecipe,
-               VPWidenStoreRecipe>(&R))
+      // Convert an unmasked or header masked scatter with an uniform address
+      // into extract-last-lane + scalar store.
+      auto *WidenStoreR = dyn_cast<VPWidenStoreRecipe>(&R);
+      if (!WidenStoreR ||
+          !vputils::isUniformAcrossVFsAndUFs(WidenStoreR->getAddr()) ||
+          WidenStoreR->isConsecutive())
         continue;
-      auto *RepR = dyn_cast<VPReplicateRecipe>(&R);
-      if (RepR && (RepR->isSingleScalar() || RepR->isPredicated()))
+      VPValue *Mask = WidenStoreR->getMask();
+
+      // Convert the scatter to a scalar store if it is unmasked or header
+      // masked.
+      if (Mask && !vputils::isHeaderMask(Mask, Plan))
         continue;
 
-      // Convert an unmasked scatter with an uniform address into
-      // extract-last-lane + scalar store.
-      // TODO: Add a profitability check comparing the cost of a scatter vs.
-      // extract + scalar store.
-      auto *WidenStoreR = dyn_cast<VPWidenStoreRecipe>(&R);
-      if (WidenStoreR && vputils::isSingleScalar(WidenStoreR->getAddr()) &&
-          !WidenStoreR->isConsecutive()) {
-        VPValue *Mask = WidenStoreR->getMask();
-
-        // Only convert the scatter to a scalar store if it is unmasked.
-        // TODO: Support converting scatter masked by the header mask to scalar
-        // store.
-        if (Mask)
+      VPInstruction *Extract;
+      if (!Mask) {
+        Extract = new VPInstruction(VPInstruction::ExtractLastLane,
+                                    {WidenStoreR->getOperand(1)});
+      } else {
+        // If the mask is the header mask, this mask contains at least one
+        // active lane. So it is safe to convert the scatter to a scalar
+        // store. Note that this will generate LastActiveLane which can only be
+        // used on header mask.
+        if (!LoopVectorizationPlanner::getDecisionAndClampRange(
+                [&](ElementCount VF) {
+                  InstructionCost ScatterCost =
+                      WidenStoreR->computeCost(VF, Ctx);
+                  // ConvertToScalarCost = LastActiveLane + ExtractLane +
+                  // scalar store.
+                  InstructionCost ScalarCost = 0;
+                  auto *ValTy =
+                      Ctx.Types.inferScalarType(WidenStoreR->getStoredValue());
+
+                  // LastActiveLane which will lower to `EVL - 1` is cheaper
+                  // under EVL.
+                  if (FoldTailWithEVL)
----------------
ElvisWang123 wrote:

Updated by using `ElementCount::getFixed(1)` as VF to query the scalar cost. Thanks!

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