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

[Florian Hahn via llvm-commits]

================
@@ -1733,6 +1733,77 @@ getUnmaskedDivRemOpcode(Intrinsic::ID ID) {
   }
 }
 
+void VPlanTransforms::narrowScatters(VPlan &Plan, VPCostContext &Ctx,
+                                     VFRange &Range, bool FoldTailWithEVL) {
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           vp_depth_first_shallow(Plan.getVectorLoopRegion()->getEntry()))) {
+    for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
+      // Convert an unmasked or header masked scatter with a uniform address
+      // into extract-last-lane + scalar store.
+      auto *WidenStoreR = dyn_cast<VPWidenStoreRecipe>(&R);
+      if (!WidenStoreR ||
+          !vputils::isUniformAcrossVFsAndUFs(WidenStoreR->getAddr()) ||
+          WidenStoreR->isConsecutive())
+        continue;
+
+      // Convert the scatter to a scalar store if it is header masked.
+      VPValue *Mask = WidenStoreR->getMask();
+      if (!Mask || !match(Mask, m_HeaderMask()))
+        continue;
+
+      // If the body is header-masked, it guarantees each iteration has at
+      // least one active lane. So it is safe to convert the scatter to a
+      // scalar store.
+      if (!LoopVectorizationPlanner::getDecisionAndClampRange(
+              [&](ElementCount VF) {
+                InstructionCost ScatterCost = WidenStoreR->computeCost(VF, Ctx);
+                Type *ValTy = WidenStoreR->getStoredValue()->getScalarType();
+                Type *AddrTy = WidenStoreR->getAddr()->getScalarType();
+                Type *IndexTy = Plan.getDataLayout().getIndexType(AddrTy);
+
+                // ScalarCost = LastActiveLaneCost + ExtractLaneCost +
+                // ScalarStoreCost.
+                InstructionCost ScalarCost = 0;
+
+                // LastActiveLane can lower to EVL - 1 when folding tail with
+                // EVL.
+                ScalarCost +=
+                    FoldTailWithEVL
+                        ? VPInstruction::getCostForRecipeWithOpcodeAndTypes(
+                              Instruction::Sub, Type::getInt32Ty(Ctx.LLVMCtx),
+                              {}, nullptr, ElementCount::getFixed(1), Ctx) +
+                              VPInstruction::getCostForRecipeWithOpcodeAndTypes(
+                                  Instruction::ZExt, IndexTy,
+                                  {Type::getInt32Ty(Ctx.LLVMCtx)}, nullptr, VF,
----------------
fhahn wrote:

Is this intentionally cost of a vector extend? I'd expect it to be a scalar extend?

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