[llvm] [LoopVectorizer] Add support for partial reductions (PR #92418)

[Graham Hunter via llvm-commits]

================
@@ -8671,6 +8754,59 @@ VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
   return tryToWiden(Instr, Operands, VPBB);
 }
 
+void VPRecipeBuilder::removeInvalidScaledReductionExitInstrs() {
+  // A partial reduction is invalid if any of its extends are used by
+  // something that isn't another partial reduction. This is because the
+  // extends are intended to be lowered along with the reduction itself.
+
+  // Build up a set of partial reduction bin ops for efficient use checking
+  SmallSet<Instruction *, 4> PartialReductionBinOps;
+  for (auto It : ScaledReductionExitInstrs) {
+    if (It.second.BinOp)
+      PartialReductionBinOps.insert(It.second.BinOp);
+  }
+
+  auto ExtendIsOnlyUsedByPartialReductions =
+      [PartialReductionBinOps](Instruction *Extend) {
+        for (auto *Use : Extend->users()) {
+          Instruction *UseInstr = dyn_cast<Instruction>(Use);
+          if (!PartialReductionBinOps.contains(UseInstr))
+            return false;
+        }
+        return true;
+      };
+
+  // Check if each use of a chain's two extends is a partial reduction
+  // and remove those that have non-partial reduction users
+  SmallSet<Instruction *, 4> PartialReductionsToRemove;
+  for (auto It : ScaledReductionExitInstrs) {
----------------
huntergr-arm wrote:

```suggestion
  for (const auto &[_, Chain] : ScaledReductionExitInstrs)
    if (!ExtendIsOnlyUsedByPartialReductions(Chain.ExtendA) ||
        !ExtendIsOnlyUsedByPartialReductions(Chain.ExtendB))
      PartialReductionsToRemove.insert(Chain.Reduction);
```

I tried finding a way to use `erase_if` for this to avoid creating a second SmallSet, but it appears we don't have a version of that which works with DenseMap yet.

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