[llvm] [VPlan] Permit more users in narrowToSingleScalars (PR #166559)

[Florian Hahn via llvm-commits]

================
@@ -1419,10 +1419,22 @@ static void narrowToSingleScalarRecipes(VPlan &Plan) {
       // broadcasts.
       if (!vputils::isSingleScalar(RepOrWidenR) ||
           !all_of(RepOrWidenR->users(), [RepOrWidenR](const VPUser *U) {
-            return U->usesScalars(RepOrWidenR) ||
-                   match(cast<VPRecipeBase>(U),
-                         m_CombineOr(m_ExtractLastElement(m_VPValue()),
-                                     m_ExtractLastLanePerPart(m_VPValue())));
+            if (auto *Store = dyn_cast<VPWidenStoreRecipe>(U)) {
+              // The assert must hold as we checked the RepOrWidenR operand
+              // against vputils::isSingleScalar.
+              assert(RepOrWidenR == Store->getAddr() ||
+                     vputils::isSingleScalar(Store->getStoredValue()));
+              return true;
+            }
----------------
fhahn wrote:

Would be good to update the comment above, to explain the store special case. From the tests is looks always beneficial. Wondering if there still could be cases where we create more broadcasts, for something like


```
%a = WIDEN add %x, 1
%b = WIDEN mul %a, 3
store %a
store %b
```

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