!fixup deep-clone whole bundle

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2697,15 +2697,18 @@ 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.
+/// 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 *> BundledOps;
+  SmallVector<VPSingleDefRecipe *> BundledRecipes;
 
   /// Temporary VPValues used for external operands of the bundle, i.e. operands
   /// not defined by recipes in the bundle.
@@ -2714,69 +2717,39 @@ class VPBundleRecipe : public VPSingleDefRecipe {
   /// Type of the bundle.
   BundleTypes BundleType;
 
-  VPBundleRecipe(BundleTypes BundleType, ArrayRef<VPSingleDefRecipe *> ToBundle)
-      : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}), BundledOps(ToBundle),
+  VPBundleRecipe(BundleTypes BundleType, ArrayRef<VPSingleDefRecipe *> ToBundle,
+                 ArrayRef<VPValue *> Operands)
+      : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}), BundledRecipes(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();
-
-    // 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 : 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());
-      }
-    }
+    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(BundleTypes::ExtendedReduction, {Ext, Red}, {}) {}
   VPBundleRecipe(VPWidenRecipe *Mul, VPReductionRecipe *Red)
-      : VPBundleRecipe(BundleTypes::MulAccumulateReduction, {Mul, Red}) {}
+      : VPBundleRecipe(BundleTypes::MulAccumulateReduction, {Mul, Red}, {}) {}
   VPBundleRecipe(VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1,
                  VPWidenRecipe *Mul, VPReductionRecipe *Red)
       : VPBundleRecipe(BundleTypes::MulAccumulateReduction,
-                       {Ext0, Ext1, Mul, Red}) {}
+                       {Ext0, Ext1, Mul, Red}, {}) {}
   VPBundleRecipe(VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1,
                  VPWidenRecipe *Mul, VPWidenCastRecipe *Ext2,
                  VPReductionRecipe *Red)
       : VPBundleRecipe(BundleTypes::MulAccumulateReduction,
-                       {Ext0, Ext1, Mul, Ext2, Red}) {}
+                       {Ext0, Ext1, Mul, Ext2, Red}, {}) {}
 
   ~VPBundleRecipe() override {
     SmallPtrSet<VPRecipeBase *, 4> Seen;
-    for (auto *R : reverse(BundledOps))
+    for (auto *R : reverse(BundledRecipes))
       if (Seen.insert(R).second)
         delete R;
     for (VPValue *T : TmpValues)
@@ -2786,13 +2759,21 @@ class VPBundleRecipe : public VPSingleDefRecipe {
   VP_CLASSOF_IMPL(VPDef::VPBundleSC)
 
   VPBundleRecipe *clone() override {
-    assert(!BundledOps.empty() && "empty bundles should be removed");
-    return new VPBundleRecipe(BundleType, BundledOps);
+    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 *getResultOp() const { return BundledOps.back(); }
+  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

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -269,8 +269,9 @@ Type *VPTypeAnalysis::inferScalarType(const VPValue *V) {
       TypeSwitch<const VPRecipeBase *, Type *>(V->getDefiningRecipe())
           .Case<VPBundleRecipe>([this](const auto *R) {
             unsigned RdxOpIdxOffset =
-                cast<VPReductionRecipe>(R->getResultOp())->isConditional() ? 2
-                                                                           : 1;
+                cast<VPReductionRecipe>(R->getResultRecipe())->isConditional()
+                    ? 2
+                    : 1;
             return inferScalarType(
                 R->getOperand(R->getNumOperands() - RdxOpIdxOffset));
           })

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2440,24 +2440,74 @@ InstructionCost VPReductionRecipe::computeCost(ElementCount VF,
                                             Ctx.CostKind);
 }
 
+void VPBundleRecipe::bundle(ArrayRef<VPValue *> Operands) {
+  assert(!BundledRecipes.empty() && "Nothing to bundle?");
+
+  // Bundle up the operand recipes.
+  SmallPtrSet<VPUser *, 4> BundledUsers;
+  for (auto *R : BundledRecipes)
+    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 (unsigned Idx = 0; Idx != BundledRecipes.size() - 1; ++Idx) {
+    VPSingleDefRecipe *R = BundledRecipes[Idx];
+    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();
+    BundledRecipes[Idx] = Copy;
+    BundledUsers.insert(Copy);
+    R->replaceUsesWithIf(Copy, [&BundledUsers](VPUser &U, unsigned) {
+      return BundledUsers.contains(&U);
+    });
+  }
+  if (BundledRecipes.back()->getParent())
+    BundledRecipes.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 : BundledRecipes) {
+    for (const auto &[Idx, Op] : enumerate(R->operands())) {
+      auto *Def = Op->getDefiningRecipe();
+      if (Def && BundledUsers.contains(Def))
+        continue;
+      if (Operands.empty())
+        addOperand(Op);
+      else
+        addOperand(Operands[TmpValues.size()]);
+      TmpValues.push_back(new VPValue());
+      R->setOperand(Idx, TmpValues.back());
+    }
+  }
+}
+
 void VPBundleRecipe::unbundle() {
-  for (auto *Op : BundledOps)
-    if (!Op->getParent())
-      Op->insertBefore(this);
+  for (auto *R : BundledRecipes)
+    if (!R->getParent())
+      R->insertBefore(this);
 
   for (const auto &[Idx, Op] : enumerate(operands()))
     TmpValues[Idx]->replaceAllUsesWith(Op);
 
-  replaceAllUsesWith(getResultOp());
+  replaceAllUsesWith(getResultRecipe());
 
   if (BundleType == BundleTypes::MulAccumulateReduction &&
-      BundledOps.size() == 5) {
+      BundledRecipes.size() == 5) {
     // Note that we will drop the extend after mul which transforms
     // reduce.add(ext(mul(ext, ext))) to reduce.add(mul(ext, ext)).
     // TODO: This transform should be done separately from bundling/unbundling.
-    auto *Ext0 = cast<VPWidenCastRecipe>(BundledOps[0]);
-    auto *Ext1 = cast<VPWidenCastRecipe>(BundledOps[1]);
-    auto *Ext2 = cast<VPWidenCastRecipe>(BundledOps[3]);
+    auto *Ext0 = cast<VPWidenCastRecipe>(BundledRecipes[0]);
+    auto *Ext1 = cast<VPWidenCastRecipe>(BundledRecipes[1]);
+    auto *Ext2 = cast<VPWidenCastRecipe>(BundledRecipes[3]);
     auto *Op0 =
         new VPWidenCastRecipe(Ext0->getOpcode(), Ext0->getOperand(0),
                               Ext2->getResultType(), *Ext0, getDebugLoc());
@@ -2469,8 +2519,8 @@ void VPBundleRecipe::unbundle() {
                                   Ext2->getResultType(), *Ext1, getDebugLoc());
       Op1->insertBefore(Ext1);
     }
-    auto *Mul = cast<VPWidenRecipe>(BundledOps[2]);
-    auto *Red = cast<VPReductionRecipe>(BundledOps[4]);
+    auto *Mul = cast<VPWidenRecipe>(BundledRecipes[2]);
+    auto *Red = cast<VPReductionRecipe>(BundledRecipes[4]);
     Mul->setOperand(0, Op0);
     Mul->setOperand(1, Op1);
     Red->setOperand(1, Mul);
@@ -2479,7 +2529,7 @@ void VPBundleRecipe::unbundle() {
     if (Ext0 != Ext1)
       Ext1->eraseFromParent();
   }
-  BundledOps.clear();
+  BundledRecipes.clear();
 }
 
 InstructionCost VPBundleRecipe::computeCost(ElementCount VF,
@@ -2492,17 +2542,17 @@ InstructionCost VPBundleRecipe::computeCost(ElementCount VF,
   switch (BundleType) {
   case BundleTypes::ExtendedReduction: {
     unsigned Opcode = RecurrenceDescriptor::getOpcode(
-        cast<VPReductionRecipe>(BundledOps[1])->getRecurrenceKind());
+        cast<VPReductionRecipe>(BundledRecipes[1])->getRecurrenceKind());
     return Ctx.TTI.getExtendedReductionCost(
         Opcode,
-        cast<VPWidenCastRecipe>(BundledOps.front())->getOpcode() ==
+        cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
             Instruction::ZExt,
         RedTy, SrcVecTy, std::nullopt, Ctx.CostKind);
   }
   case BundleTypes::MulAccumulateReduction:
     return Ctx.TTI.getMulAccReductionCost(
-        BundledOps.size() > 2
-            ? cast<VPWidenCastRecipe>(BundledOps.front())->getOpcode() ==
+        BundledRecipes.size() > 2
+            ? cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
                   Instruction::ZExt
             : false,
         RedTy, SrcVecTy, Ctx.CostKind);
@@ -2516,7 +2566,7 @@ void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
   O << Indent << "BUNDLE ";
   printAsOperand(O, SlotTracker);
   O << " = ";
-  auto *Red = cast<VPReductionRecipe>(BundledOps.back());
+  auto *Red = cast<VPReductionRecipe>(BundledRecipes.back());
   unsigned Opcode = RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind());
 
   switch (BundleType) {
@@ -2527,7 +2577,7 @@ void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
     getOperand(0)->printAsOperand(O, SlotTracker);
     Red->printFlags(O);
 
-    auto *Ext0 = cast<VPWidenCastRecipe>(BundledOps[0]);
+    auto *Ext0 = cast<VPWidenCastRecipe>(BundledRecipes[0]);
     O << Instruction::getOpcodeName(Ext0->getOpcode()) << " to "
       << *Ext0->getResultType();
     if (Red->isConditional()) {
@@ -2545,16 +2595,16 @@ void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
              RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()))
       << " (";
     O << "mul";
-    auto *Mul = cast<VPWidenRecipe>(BundledOps.size() == 2 ? BundledOps[0]
-                                                           : BundledOps[2]);
+    auto *Mul = cast<VPWidenRecipe>(
+        BundledRecipes.size() == 2 ? BundledRecipes[0] : BundledRecipes[2]);
     Mul->printFlags(O);
-    bool IsExtended = BundledOps.size() > 2;
+    bool IsExtended = BundledRecipes.size() > 2;
     if (IsExtended)
       O << "(";
     getOperand(0)->printAsOperand(O, SlotTracker);
     if (IsExtended) {
       auto *Ext0 = cast<VPWidenCastRecipe>(
-          BundledOps.size() == 5 ? BundledOps[3] : BundledOps[0]);
+          BundledRecipes.size() == 5 ? BundledRecipes[3] : BundledRecipes[0]);
       O << " " << Instruction::getOpcodeName(Ext0->getOpcode()) << " to "
         << *Ext0->getResultType() << "), (";
     } else {
@@ -2563,7 +2613,7 @@ void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
     getOperand(1)->printAsOperand(O, SlotTracker);
     if (IsExtended) {
       auto *Ext1 = cast<VPWidenCastRecipe>(
-          BundledOps.size() == 5 ? BundledOps[3] : BundledOps[1]);
+          BundledRecipes.size() == 5 ? BundledRecipes[3] : BundledRecipes[1]);
       O << " " << Instruction::getOpcodeName(Ext1->getOpcode()) << " to "
         << *Ext1->getResultType() << ")";
     }