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

[Florian Hahn via llvm-commits]

================
@@ -2580,30 +2580,174 @@ InstructionCost VPReductionRecipe::computeCost(ElementCount VF,
                                             Ctx.CostKind);
 }
 
-InstructionCost
-VPExtendedReductionRecipe::computeCost(ElementCount VF,
-                                       VPCostContext &Ctx) const {
-  unsigned Opcode = RecurrenceDescriptor::getOpcode(getRecurrenceKind());
+VPSingleDefBundleRecipe::VPSingleDefBundleRecipe(
+    BundleTypes BundleType, ArrayRef<VPSingleDefRecipe *> ToBundle)
+    : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}),
+      BundledRecipes(
+          SetVector<VPSingleDefRecipe *>(ToBundle.begin(), ToBundle.end())
+              .takeVector()),
+      BundleType(BundleType) {
+  assert(!BundledRecipes.empty() && "Nothing to bundle?");
+
+  // Maintain a copy of the bundled recipes as a set of users.
+  SmallPtrSet<VPUser *, 4> BundledRecipesAsSetOfUsers;
+  for (auto *R : BundledRecipes)
+    BundledRecipesAsSetOfUsers.insert(R);
+
+  // Recipes in the bundle, except the last one, must only be used by (other)
+  // recipes inside the bundle. If there are other users, external to the
+  // bundle, use a clone of the recipe for external users.
+  for (VPSingleDefRecipe *R : BundledRecipes) {
+    if (R != BundledRecipes.back() &&
+        any_of(R->users(), [&BundledRecipesAsSetOfUsers](VPUser *U) {
+          return !BundledRecipesAsSetOfUsers.contains(U);
+        })) {
+      // There are users outside of the bundle. Clone the recipe and use the
+      // clone those external users.
+      VPSingleDefRecipe *CopyForExtUsers = R->clone();
+      R->replaceUsesWithIf(CopyForExtUsers,
+                           [&BundledRecipesAsSetOfUsers](VPUser &U, unsigned) {
+                             return !BundledRecipesAsSetOfUsers.contains(&U);
+                           });
+      CopyForExtUsers->insertBefore(R);
+    }
+    if (R->getParent())
+      R->removeFromParent();
+  }
+
+  // Internalize all external operands to the bundled recipes. To do so,
+  // create new temporary VPValues for all operands defined by a recipe outside
+  // the bundle. The original operands are added as operands of the
+  // VPSingleDefBundleRecipe itself.
+  for (auto *R : BundledRecipes) {
+    for (const auto &[Idx, Op] : enumerate(R->operands())) {
+      auto *Def = Op->getDefiningRecipe();
+      if (Def && BundledRecipesAsSetOfUsers.contains(Def))
+        continue;
+      addOperand(Op);
+      BundleLiveInPlaceholders.push_back(new VPValue());
+      R->setOperand(Idx, BundleLiveInPlaceholders.back());
+    }
+  }
+}
+
+void VPSingleDefBundleRecipe::unbundle() {
+  for (auto *R : BundledRecipes)
+    R->insertBefore(this);
+
+  for (const auto &[Idx, Op] : enumerate(operands()))
+    BundleLiveInPlaceholders[Idx]->replaceAllUsesWith(Op);
+
+  replaceAllUsesWith(BundledRecipes.back());
+  BundledRecipes.clear();
+}
+
+InstructionCost VPSingleDefBundleRecipe::computeCost(ElementCount VF,
+                                                     VPCostContext &Ctx) const {
   Type *RedTy = Ctx.Types.inferScalarType(this);
-  auto *SrcVecTy =
-      cast<VectorType>(toVectorTy(Ctx.Types.inferScalarType(getVecOp()), VF));
+  auto *SrcVecTy = cast<VectorType>(
+      toVectorTy(Ctx.Types.inferScalarType(getOperand(0)), VF));
   assert(RedTy->isIntegerTy() &&
-         "ExtendedReduction only support integer type currently.");
-  return Ctx.TTI.getExtendedReductionCost(Opcode, isZExt(), RedTy, SrcVecTy,
-                                          std::nullopt, Ctx.CostKind);
+         "VPSingleDefBundleRecipe only supports integer types currently.");
+  switch (BundleType) {
+  case BundleTypes::ExtendedReduction: {
+    unsigned Opcode = RecurrenceDescriptor::getOpcode(
+        cast<VPReductionRecipe>(BundledRecipes[1])->getRecurrenceKind());
+    return Ctx.TTI.getExtendedReductionCost(
+        Opcode,
+        cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
+            Instruction::ZExt,
+        RedTy, SrcVecTy, std::nullopt, Ctx.CostKind);
+  }
+  case BundleTypes::MulAccumulateReduction:
+    return Ctx.TTI.getMulAccReductionCost(false, RedTy, SrcVecTy, Ctx.CostKind);
+
+  case BundleTypes::ExtMulAccumulateReduction:
+    return Ctx.TTI.getMulAccReductionCost(
+        cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
+            Instruction::ZExt,
+        RedTy, SrcVecTy, Ctx.CostKind);
+  }
+}
+
+bool VPSingleDefBundleRecipe::mayReadOrWriteMemory() const {
+  return any_of(BundledRecipes, [](VPSingleDefRecipe *R) {
+    return R->mayReadFromMemory() || R->mayWriteToMemory();
+  });
 }
 
-InstructionCost
-VPMulAccumulateReductionRecipe::computeCost(ElementCount VF,
-                                            VPCostContext &Ctx) const {
-  Type *RedTy = Ctx.Types.inferScalarType(this);
-  auto *SrcVecTy =
-      cast<VectorType>(toVectorTy(Ctx.Types.inferScalarType(getVecOp0()), VF));
-  return Ctx.TTI.getMulAccReductionCost(isZExt(), RedTy, SrcVecTy,
-                                        Ctx.CostKind);
+bool VPSingleDefBundleRecipe::mayHaveSideEffects() const {
+  return any_of(BundledRecipes,
+                [](VPSingleDefRecipe *R) { return R->mayHaveSideEffects(); });
----------------
fhahn wrote:

For now, I added an assert here and. the constructor to disallow recieps with side-effects, thanks

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