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

[Florian Hahn via llvm-commits]

================
@@ -2917,6 +2706,108 @@ 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. The bundled recipes are completely connected from the def-use graph
+/// outside the bundled recipes. Operands not defined by recipes in the bundle
+/// are added as operands of the VPBundleRecipe and the users of the result
+/// recipe must be updated to use the VPBundleRecipe.
+class VPBundleRecipe : public VPSingleDefRecipe {
+  enum class BundleTypes {
+    ExtendedReduction,
+    MulAccumulateReduction,
+  };
+
+  /// Recipes bundled together in this VPBundleRecipe.
+  SmallVector<VPSingleDefRecipe *> BundledRecipes;
+
+  /// 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,
+                 ArrayRef<VPValue *> Operands)
+      : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}), BundledRecipes(ToBundle),
+        BundleType(BundleType) {
+    bundle(Operands);
+  }
+
+  /// Internalize recipes in BundledRecipes External operands (i.e. not defined
+  /// by another recipe in the bundle) are replaced by temporary VPValues and
+  /// the original operands are transferred to the VPBundleRecipe itself. Clone
+  /// recipes as needed to ensure they are only used by other recipes in the
+  /// bundle. If \p Operands is not empty, use it as operands for the new
+  /// VPBundleRecipe (used when cloning the recipe).
+  void bundle(ArrayRef<VPValue *> Operands);
+
+public:
+  VPBundleRecipe(VPWidenCastRecipe *Ext, VPReductionRecipe *Red)
+      : VPBundleRecipe(BundleTypes::ExtendedReduction, {Ext, Red}, {}) {}
+  VPBundleRecipe(VPWidenRecipe *Mul, VPReductionRecipe *Red)
+      : VPBundleRecipe(BundleTypes::MulAccumulateReduction, {Mul, Red}, {}) {}
+  VPBundleRecipe(VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1,
+                 VPWidenRecipe *Mul, VPReductionRecipe *Red)
+      : VPBundleRecipe(BundleTypes::MulAccumulateReduction,
+                       {Ext0, Ext1, Mul, Red}, {}) {}
+  VPBundleRecipe(VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1,
+                 VPWidenRecipe *Mul, VPWidenCastRecipe *Ext2,
+                 VPReductionRecipe *Red)
+      : VPBundleRecipe(BundleTypes::MulAccumulateReduction,
+                       {Ext0, Ext1, Mul, Ext2, Red}, {}) {}
+
+  ~VPBundleRecipe() override {
+    SmallPtrSet<VPRecipeBase *, 4> Seen;
+    for (auto *R : reverse(BundledRecipes))
+      if (Seen.insert(R).second)
+        delete R;
+    for (VPValue *T : TmpValues)
+      delete T;
+  }
+
+  VP_CLASSOF_IMPL(VPDef::VPBundleSC)
+
+  VPBundleRecipe *clone() override {
+    assert(!BundledRecipes.empty() && "empty bundles should be removed");
+    SmallVector<VPSingleDefRecipe *> NewBundledRecipes;
+    for (auto *R : BundledRecipes)
+      NewBundledRecipes.push_back(R->clone());
+    for (auto *New : NewBundledRecipes) {
+      for (const auto &[Idx, Old] : enumerate(BundledRecipes)) {
+        New->replaceUsesOfWith(Old, NewBundledRecipes[Idx]);
+      }
+    }
+    return new VPBundleRecipe(BundleType, NewBundledRecipes, operands());
+  }
+
+  /// Return the VPSingleDefRecipe producing the final result of the bundled
+  /// recipe.
+  VPSingleDefRecipe *getResultRecipe() const { return BundledRecipes.back(); }
+
+  /// Insert the bundled recipes back into the VPlan, directly before the
+  /// current recipe. Leaves the bundle recipe empty and the recipe must be
+  /// removed before codegen.
+  void unbundle();
+
+  /// Generate the extraction of the appropriate bit from the block mask and the
+  /// conditional branch.
----------------
fhahn wrote:

Updated, thanks

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