[llvm] [VPlan] Add VPBundleRecipe, replacing extended reduction recipes. (PR #144281)

[Sam Tebbs via llvm-commits]

================
@@ -2922,6 +2711,127 @@ class VPBranchOnMaskRecipe : public VPRecipeBase {
   }
 };
 
+/// A recipe to combine multiple recipes into a 'bundle' recipe, which should be
+/// considered as single entity for cost-modeling and transforms. The recipe
+/// needs to be 'unbundled', i.e. replaced by its individual recipes before
+/// execute.
+class VPBundleRecipe : public VPSingleDefRecipe {
+  enum class BundleTypes {
+    ExtendedReduction,
+    MulAccumulateReduction,
+  };
+
+  /// Recipes bundled together in this VPBundleRecipe.
+  SmallVector<VPSingleDefRecipe *> BundledOps;
+
+  /// Temporary VPValues used for external operands of the bundle, i.e. operands
+  /// not defined by recipes in the bundle.
+  SmallVector<VPValue *> TmpValues;
+
+  /// Type of the bundle.
+  BundleTypes BundleType;
+
+  VPBundleRecipe(BundleTypes BundleType, ArrayRef<VPSingleDefRecipe *> ToBundle)
+      : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}), BundledOps(ToBundle),
+        BundleType(BundleType) {
+    // Bundle up the operand recipes.
+    SmallPtrSet<VPUser *, 4> BundledUsers;
+    for (auto *R : BundledOps)
+      BundledUsers.insert(R);
+
+    // Recipes in the bundle, except the last one, must only be used inside the
+    // bundle. If there other external users, clone the recipes for the bundle.
+    for (const auto &[Idx, R] : enumerate(drop_end(ToBundle))) {
+      if (all_of(R->users(), [&BundledUsers](VPUser *U) {
+            return BundledUsers.contains(U);
+          })) {
+        if (R->getParent())
+          R->removeFromParent();
+        continue;
+      }
+      // The users external to the bundle. Clone the recipe for use in the
+      // bundle and update all its in-bundle users.
+      VPSingleDefRecipe *Copy = R->clone();
+      BundledOps[Idx] = Copy;
+      BundledUsers.insert(Copy);
+      R->replaceUsesWithIf(Copy, [&BundledUsers](VPUser &U, unsigned) {
+        return BundledUsers.contains(&U);
+      });
+    }
+    BundledOps.back()->removeFromParent();
----------------
SamTebbs33 wrote:

While experimenting with rebasing https://github.com/llvm/llvm-project/pull/136173 on top of this, I've found that the new bundle recipe can't be used when interleaving. This line for example needs to be gated behind `if (BundledOps.back().getParent())` since the second time the bundle is created for the same reduction recipe, it's already been removed and this causes an assertion failure. Would you mind seeing if you can get this recipe to work for interleaved reductions?

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