[llvm] [LV] Check all users of partial reductions in chain have same scale. (PR #162822)

[Florian Hahn via llvm-commits]

================
@@ -7910,6 +7910,29 @@ void VPRecipeBuilder::collectScaledReductions(VFRange &Range) {
         (!Chain.ExtendB || ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB)))
       ScaledReductionMap.try_emplace(Chain.Reduction, Pair.second);
   }
+
+  // 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.
+  auto AllUsersPartialRdx = [this](Instruction *I, unsigned Scale) {
+    return all_of(I->users(), [Scale, this](const User *U) {
+      auto *UI = cast<Instruction>(U);
+
+      if (isa<PHINode>(UI) && UI->getParent() == OrigLoop->getHeader()) {
+        return all_of(UI->users(), [Scale, this](const User *U) {
+          auto *UI = cast<Instruction>(U);
+          return ScaledReductionMap.lookup_or(UI, 0) == Scale;
+        });
+      }
+
+      return ScaledReductionMap.lookup_or(UI, 0) == Scale ||
+             !OrigLoop->contains(UI->getParent());
+    });
+  };
+  for (const auto &[Chain, Scale] : PartialReductionChains) {
+    if (!AllUsersPartialRdx(Chain.Reduction, Scale))
+      ScaledReductionMap.erase(Chain.Reduction);
----------------
fhahn wrote:

We could use the information from `PartialReductionChains`, but this would include entries that will get rejected later, by the `ExtendIsOnlyUsedByPartialReductions`. I kept it as 2 separate loops for now.

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