!fixup address latest comments, thanks

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2707,93 +2707,95 @@ 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.
+/// considered a single entity for cost-modeling and transforms. The recipe
+/// needs to be 'unbundled', i.e. replaced by its bundled recipes before
+/// execute. The bundled recipes are completely disconnected from the def-use
+/// graph of other, non-bundled recipes. Def-use edges between pairs of bundled
+/// recipes remain intact, whereas every edge between a bundled and a
+/// non-bundled recipe is elevated to connect the non-bundled recipe with the
+/// VPExpression itself.
+class VPExpression : public VPSingleDefRecipe {
+  /// Recipes bundled together in this VPExpression.
   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;
+  SmallVector<VPValue *> BundleLiveInPlaceholders;
+
+  enum class BundleTypes {
+    /// Represents an inloop extended reduction operation, performing a
+    /// reduction on a extended vector operand into a scalar value, and adding
+    /// the result to a chain.
+    ExtendedReduction,
+    /// Represent an inloop multiply-accumulate reduction, multiplying the
+    /// extended vector operands, performing a reduction.add on the result, and
+    /// adding the scalar result to a chain.
+    ExtMulAccumulateReduction,
+    /// Represent an inloop multiply-accumulate reduction, multiplying the
+    /// vector operands, performing a reduction.add on the result, and adding
+    /// the scalar result to a chain.
+    MulAccumulateReduction,
+  };
 
   /// 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);
+  /// Construct a new VPExpression by internalizing recipes in \p
+  /// 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 VPExpression itself. Clone recipes as needed
+  /// (excluding last) to ensure they are only used by other recipes in the
+  /// bundle.
+  VPExpression(BundleTypes BundleType, ArrayRef<VPSingleDefRecipe *> ToBundle);
 
 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 {
+  VPExpression(VPWidenCastRecipe *Ext, VPReductionRecipe *Red)
+      : VPExpression(BundleTypes::ExtendedReduction, {Ext, Red}) {}
+  VPExpression(VPWidenRecipe *Mul, VPReductionRecipe *Red)
+      : VPExpression(BundleTypes::MulAccumulateReduction, {Mul, Red}) {}
+  VPExpression(VPWidenCastRecipe *Ext0, VPWidenCastRecipe *Ext1,
+               VPWidenRecipe *Mul, VPReductionRecipe *Red)
+      : VPExpression(BundleTypes::ExtMulAccumulateReduction,
+                     {Ext0, Ext1, Mul, Red}) {}
+
+  ~VPExpression() override {
     SmallPtrSet<VPRecipeBase *, 4> Seen;
     for (auto *R : reverse(BundledRecipes))
       if (Seen.insert(R).second)
         delete R;
-    for (VPValue *T : TmpValues)
+    for (VPValue *T : BundleLiveInPlaceholders)
       delete T;
   }
 
   VP_CLASSOF_IMPL(VPDef::VPBundleSC)
 
-  VPBundleRecipe *clone() override {
+  VPExpression *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)) {
+      for (const auto &[Idx, Old] : enumerate(BundledRecipes))
         New->replaceUsesOfWith(Old, NewBundledRecipes[Idx]);
-      }
+      // Update placeholder operands in the cloned recipe to use the external
+      // operands, to be internalized when the cloned bundle is constructed.
+      for (const auto &[Placeholder, OutsideOp] :
+           zip(BundleLiveInPlaceholders, operands()))
+        New->replaceUsesOfWith(Placeholder, OutsideOp);
     }
-    return new VPBundleRecipe(BundleType, NewBundledRecipes, operands());
+    return new VPExpression(BundleType, NewBundledRecipes);
   }
 
   /// 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.
+  /// current recipe. Leaves the bundle recipe empty, which must be removed
+  /// before codegen.
   void unbundle();
 
-  /// Generate the extraction of the appropriate bit from the block mask and the
-  /// conditional branch.
+  /// Method for generating code, must not be called as this recipe is abstract.
   void execute(VPTransformState &State) override {
     llvm_unreachable("recipe must be removed before execute");
   }
@@ -2806,6 +2808,13 @@ class VPBundleRecipe : public VPSingleDefRecipe {
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
+
+  /// Returns true if this bundle contains recipes that may read from or write
+  /// to memory.
+  bool mayReadOrWriteMemory() const;
+
+  /// Returns true if this bundle contains recipes that may have side effects.
+  bool mayHaveSideEffects() const;
 };
 
 /// VPPredInstPHIRecipe is a recipe for generating the phi nodes needed when

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -49,6 +49,8 @@ using VectorParts = SmallVector<Value *, 2>;
 
 bool VPRecipeBase::mayWriteToMemory() const {
   switch (getVPDefID()) {
+  case VPBundleSC:
+    return cast<VPBundleRecipe>(this)->mayReadOrWriteMemory();
   case VPInstructionSC:
     return cast<VPInstruction>(this)->opcodeMayReadOrWriteFromMemory();
   case VPInterleaveSC:
@@ -97,6 +99,8 @@ bool VPRecipeBase::mayWriteToMemory() const {
 
 bool VPRecipeBase::mayReadFromMemory() const {
   switch (getVPDefID()) {
+  case VPBundleSC:
+    return cast<VPBundleRecipe>(this)->mayReadOrWriteMemory();
   case VPInstructionSC:
     return cast<VPInstruction>(this)->opcodeMayReadOrWriteFromMemory();
   case VPWidenLoadEVLSC:
@@ -143,6 +147,8 @@ bool VPRecipeBase::mayReadFromMemory() const {
 
 bool VPRecipeBase::mayHaveSideEffects() const {
   switch (getVPDefID()) {
+  case VPBundleSC:
+    return cast<VPBundleRecipe>(this)->mayHaveSideEffects();
   case VPDerivedIVSC:
   case VPFirstOrderRecurrencePHISC:
   case VPPredInstPHISC:
@@ -2500,95 +2506,65 @@ InstructionCost VPReductionRecipe::computeCost(ElementCount VF,
                                             Ctx.CostKind);
 }
 
-void VPBundleRecipe::bundle(ArrayRef<VPValue *> Operands) {
+VPBundleRecipe::VPBundleRecipe(BundleTypes BundleType,
+                               ArrayRef<VPSingleDefRecipe *> ToBundle)
+    : VPSingleDefRecipe(VPDef::VPBundleSC, {}, {}),
+      BundledRecipes(
+          SetVector<VPSingleDefRecipe *>(ToBundle.begin(), ToBundle.end())
+              .takeVector()),
+      BundleType(BundleType) {
   assert(!BundledRecipes.empty() && "Nothing to bundle?");
 
-  // Bundle up the operand recipes.
-  SmallPtrSet<VPUser *, 4> BundledUsers;
+  // Maintain a copy of the bundled recipes as a set of users.
+  SmallPtrSet<VPUser *, 4> BundledRecipesAsSetOfUsers;
   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);
+    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);
         })) {
-      if (R->getParent())
-        R->removeFromParent();
-      continue;
+      // 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);
     }
-    // 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
+    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
-  // VPBundleRecipe.
+  // VPBundleRecipe itself.
   for (auto *R : BundledRecipes) {
     for (const auto &[Idx, Op] : enumerate(R->operands())) {
       auto *Def = Op->getDefiningRecipe();
-      if (Def && BundledUsers.contains(Def))
+      if (Def && BundledRecipesAsSetOfUsers.contains(Def))
         continue;
-      if (Operands.empty())
-        addOperand(Op);
-      else
-        addOperand(Operands[TmpValues.size()]);
-      TmpValues.push_back(new VPValue());
-      R->setOperand(Idx, TmpValues.back());
+      addOperand(Op);
+      BundleLiveInPlaceholders.push_back(new VPValue());
+      R->setOperand(Idx, BundleLiveInPlaceholders.back());
     }
   }
 }
 
 void VPBundleRecipe::unbundle() {
   for (auto *R : BundledRecipes)
-    if (!R->getParent())
       R->insertBefore(this);
 
   for (const auto &[Idx, Op] : enumerate(operands()))
-    TmpValues[Idx]->replaceAllUsesWith(Op);
+    BundleLiveInPlaceholders[Idx]->replaceAllUsesWith(Op);
 
   replaceAllUsesWith(getResultRecipe());
-
-  if (BundleType == BundleTypes::MulAccumulateReduction &&
-      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>(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());
-    Op0->insertBefore(Ext0);
-
-    VPSingleDefRecipe *Op1 = Op0;
-    if (Ext0 != Ext1) {
-      Op1 = new VPWidenCastRecipe(Ext1->getOpcode(), Ext1->getOperand(0),
-                                  Ext2->getResultType(), *Ext1, getDebugLoc());
-      Op1->insertBefore(Ext1);
-    }
-    auto *Mul = cast<VPWidenRecipe>(BundledRecipes[2]);
-    auto *Red = cast<VPReductionRecipe>(BundledRecipes[4]);
-    Mul->setOperand(0, Op0);
-    Mul->setOperand(1, Op1);
-    Red->setOperand(1, Mul);
-    Ext0->eraseFromParent();
-    Ext2->eraseFromParent();
-    if (Ext0 != Ext1)
-      Ext1->eraseFromParent();
-  }
   BundledRecipes.clear();
 }
 
@@ -2610,15 +2586,27 @@ InstructionCost VPBundleRecipe::computeCost(ElementCount VF,
         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(
-        BundledRecipes.size() > 2
-            ? cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
-                  Instruction::ZExt
-            : false,
+        cast<VPWidenCastRecipe>(BundledRecipes.front())->getOpcode() ==
+            Instruction::ZExt,
         RedTy, SrcVecTy, Ctx.CostKind);
   }
 }
 
+bool VPBundleRecipe::mayReadOrWriteMemory() const {
+  return any_of(BundledRecipes, [](VPSingleDefRecipe *R) {
+    return R->mayReadFromMemory() || R->mayWriteToMemory();
+  });
+}
+
+bool VPBundleRecipe::mayHaveSideEffects() const {
+  return any_of(BundledRecipes,
+                [](VPSingleDefRecipe *R) { return R->mayHaveSideEffects(); });
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 
 void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
@@ -2647,33 +2635,32 @@ void VPBundleRecipe::print(raw_ostream &O, const Twine &Indent,
     O << ")";
     break;
   }
-  case BundleTypes::MulAccumulateReduction: {
+  case BundleTypes::MulAccumulateReduction:
+  case BundleTypes::ExtMulAccumulateReduction: {
     getOperand(getNumOperands() - 1)->printAsOperand(O, SlotTracker);
     O << " + ";
     O << "reduce."
       << Instruction::getOpcodeName(
              RecurrenceDescriptor::getOpcode(Red->getRecurrenceKind()))
       << " (";
     O << "mul";
-    auto *Mul = cast<VPWidenRecipe>(
-        BundledRecipes.size() == 2 ? BundledRecipes[0] : BundledRecipes[2]);
+    bool IsExtended = BundleType == BundleTypes::ExtMulAccumulateReduction;
+    auto *Mul =
+        cast<VPWidenRecipe>(IsExtended ? BundledRecipes[2] : BundledRecipes[0]);
     Mul->printFlags(O);
-    bool IsExtended = BundledRecipes.size() > 2;
     if (IsExtended)
       O << "(";
     getOperand(0)->printAsOperand(O, SlotTracker);
     if (IsExtended) {
-      auto *Ext0 = cast<VPWidenCastRecipe>(
-          BundledRecipes.size() == 5 ? BundledRecipes[3] : BundledRecipes[0]);
+      auto *Ext0 = cast<VPWidenCastRecipe>(BundledRecipes[0]);
       O << " " << Instruction::getOpcodeName(Ext0->getOpcode()) << " to "
         << *Ext0->getResultType() << "), (";
     } else {
       O << ", ";
     }
     getOperand(1)->printAsOperand(O, SlotTracker);
     if (IsExtended) {
-      auto *Ext1 = cast<VPWidenCastRecipe>(
-          BundledRecipes.size() == 5 ? BundledRecipes[3] : BundledRecipes[1]);
+      auto *Ext1 = cast<VPWidenCastRecipe>(BundledRecipes[1]);
       O << " " << Instruction::getOpcodeName(Ext1->getOpcode()) << " to "
         << *Ext1->getResultType() << ")";
     }