[llvm] [LV] Convert uniform-address scatters to scalar store when unmasked or header-masked. (PR #166114)

[Mel Chen via llvm-commits]

================
@@ -1372,6 +1372,50 @@ void VPlanTransforms::simplifyRecipes(VPlan &Plan) {
   }
 }
 
+static VPSingleDefRecipe *findHeaderMask(VPlan &Plan);
+
+/// Convert scatters with a uniform address that are either unmasked or
+/// masked by the header mask into an extract-last-element + scalar store.
+//  TODO: Add a profitability check comparing the cost of a scatter vs.
+//  extract + scalar store.
+static void optimizeScatterWithUniformAddr(VPlan &Plan) {
+  VPValue *HeaderMask = findHeaderMask(Plan);
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           vp_depth_first_deep(Plan.getEntry()))) {
+    for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
+
+      // Only transform store recipes.
+      if (!isa<VPWidenStoreRecipe, VPWidenStoreEVLRecipe>(&R))
+        continue;
+
+      auto StoreR = cast<VPWidenMemoryRecipe>(&R);
+      if (StoreR->isConsecutive() ||
+          !vputils::isSingleScalar(StoreR->getAddr()))
+        continue;
+
+      assert(!StoreR->isReverse() &&
+             "Not consecutive memory recipes shouldn't be reversed");
+      VPValue *Mask = StoreR->getMask();
+
+      // Only convert the scatter to a scalar store if it is unmasked or masked
+      // by the header mask, which guarantees that at least one active lane.
+      if (Mask && Mask != HeaderMask)
+        continue;
----------------
Mel-Chen wrote:

```suggestion
      // Only convert the scatter to a scalar store if it is unmasked.
      // TODO: Support header mask.
      if (StoreR->isMasked())
        continue;
```

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