[llvm] [LV] Check full partial reduction chains in order. (PR #168036)

[Florian Hahn via llvm-commits]

================
@@ -8016,31 +8019,41 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
 
   // Check if each use of a chain's two extends is a partial reduction
   // and only add those that don't have non-partial reduction users.
-  for (auto Pair : PartialReductionChains) {
-    PartialReductionChain Chain = Pair.first;
-    if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
-        (!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
-      ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
+  for (const auto &[_, Chains] : ChainsByPhi) {
+    for (const auto &[Chain, Scale] : Chains) {
+      if (ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) &&
+          (!Chain.ExtendB ||
+           ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
+        ScaledReductionMap.try_emplace(Chain.Reduction, Scale);
+    }
   }
 
   // Check that all partial reductions in a chain are only used by other
   // partial reductions with the same scale factor. Otherwise we end up creating
   // users of scaled reductions where the types of the other operands don't
   // match.
-  for (const auto &[Chain, Scale] : PartialReductionChains) {
-    auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
-      auto *UI = cast<Instruction>(U);
-      if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
-        return all_of(UI->users(), [ScaleVal, this](const User *U) {
-          auto *UI = cast<Instruction>(U);
-          return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
-        });
+  for (const auto &[Phi, Chains] : ChainsByPhi) {
+    for (const auto &[Chain, Scale] : Chains) {
+      auto AllUsersPartialRdx = [ScaleVal = Scale, this](const User *U) {
+        auto *UI = cast<Instruction>(U);
+        if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
+          return all_of(UI->users(), [ScaleVal, this](const User *U) {
+            auto *UI = cast<Instruction>(U);
+            return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal;
+          });
+        }
+        return ScaledReductionMap.lookup_or(UI, 0) == ScaleVal ||
+               !OrigLoop->contains(UI->getParent());
+      };
----------------
fhahn wrote:


`AllUsersPartialRdx` should be the same as before this patch, but we need an extra loop level.

It checks both if the chain is complete, without any gaps/missing links (all in loop users must also be in the map) and that their scale factor match. The latter could be checked by simply iterating over all entries in the map, but the former unfortunately cannot currently. `llvm/test/Transforms/LoopVectorize/AArch64/partial-reduce-incomplete-chains.ll` has a few tests with chains with missing links which need the first check.

But it is now no longer necessary to iterate over all users of phis, we can simply check if it is the phi for the chain, thanks!



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