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

[Florian Hahn via llvm-commits]

================
@@ -1845,6 +1845,91 @@ static void reassociateHeaderMask(VPlan &Plan) {
   }
 }
 
+void VPlanTransforms::narrowScatters(VPlan &Plan, VPCostContext &Ctx,
+                                     VFRange &Range,
+                                     const bool &FoldTailWithEVL) {
+  if (Plan.hasScalarVFOnly())
+    return;
+
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()))) {
+    for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
+      // 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;
+      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;
+
+      if (!Mask) {
+        VPInstruction *Extract = new VPInstruction(
+            VPInstruction::ExtractLastLane, {WidenStoreR->getOperand(1)});
+        // TODO: Sink the scalar store recipe to middle block if possible.
+        auto *ScalarStore = new VPReplicateRecipe(
+            &WidenStoreR->getIngredient(), {Extract, WidenStoreR->getAddr()},
+            true /*IsSingleScalar*/, nullptr /*Mask*/, {},
+            *WidenStoreR /*Metadata*/);
+        Extract->insertBefore(WidenStoreR);
+        ScalarStore->insertBefore(WidenStoreR);
+        WidenStoreR->eraseFromParent();
+      } else {
+        VPBuilder Builder(WidenStoreR);
+        VPInstruction *LastActiveLane =
+            Builder.createNaryOp(VPInstruction::LastActiveLane, {Mask});
+        VPInstruction *Extract = Builder.createNaryOp(
+            VPInstruction::ExtractLane,
+            {LastActiveLane, WidenStoreR->getStoredValue()});
+        auto *ScalarStore = new VPReplicateRecipe(
+            &WidenStoreR->getIngredient(), {Extract, WidenStoreR->getAddr()},
+            true /*IsSingleScalar*/, nullptr /*Mask*/, {},
+            *WidenStoreR /*Metadata*/);
+        ScalarStore->insertBefore(WidenStoreR);
+
+        // 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;
+
+                  // LastActiveLane can lower to EVL - 1 when folding tail with
+                  // EVL.
+                  if (FoldTailWithEVL)
+                    ScalarCost += Ctx.TTI.getArithmeticInstrCost(
+                        Instruction::Sub, Type::getInt32Ty(Ctx.LLVMCtx),
+                        Ctx.CostKind);
----------------
fhahn wrote:

Yes, although as @lukel97 also mentioned, it would be great if that could be avoided. I also don't have a good suggestion, other than checking if we can run this late, after we introduced EVL? Or is there a particular reason we need to run this before then?

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