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

[Sam Tebbs via llvm-commits]

================
@@ -2922,6 +2710,123 @@ 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 : ToBundle)
+      BundledUsers.insert(R);
+
+    // Recipes in the bundle, expect 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;
+      }
+      // There users external to the bundle. Clone the recipe for use in the
+      // bundle and update all its in-bundle users.
+      this->BundledOps[Idx] = R->clone();
+      BundledUsers.insert(this->BundledOps[Idx]);
+      R->replaceUsesWithIf(this->BundledOps[Idx],
+                           [&BundledUsers](VPUser &U, unsigned) {
+                             return BundledUsers.contains(&U);
+                           });
+    }
+    BundledOps.back()->removeFromParent();
+
+    // Internalize all external operands to the bundled operations. To do so,
+    // create new temporary VPValues for all operands not defined by recipe in
+    // the bundle. The original operands are added as operands of the
+    // VPBundleRecipe.
+    for (auto *R : this->BundledOps) {
+      for (const auto &[Idx, Op] : enumerate(R->operands())) {
+        auto *Def = Op->getDefiningRecipe();
+        if (Def && BundledUsers.contains(Def))
+          continue;
+        addOperand(Op);
+        TmpValues.push_back(new VPValue());
+        R->setOperand(Idx, TmpValues.back());
+      }
+    }
+  }
+
+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(BundledOps))
+      if (Seen.insert(R).second)
+        delete R;
+    for (VPValue *T : TmpValues)
+      delete T;
+  }
+
+  VP_CLASSOF_IMPL(VPDef::VPBundleSC)
+
+  VPBundleRecipe *clone() override {
+    return new VPBundleRecipe(BundleType, BundledOps);
+  }
+
+  /// Return the VPSingleDefRecipe producing the final result of the bundled
+  /// recipe.
+  VPSingleDefRecipe *getResultOp() const { return BundledOps.back(); }
+
+  void unbundle();
----------------
SamTebbs33 wrote:

Docs needed.

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