[VPlan, VP] Introduce new recipes to support predicated vectorization

[vkmr]

Summary

This patch introduces new VPlan recipes VPWidenEVLRecipe, VPPredicatedWidenRecipe, VPPredicatedWidenMemoryInstrctionRecipe, and AllTrueMask VPInstruction as a first step to enable predicated vectorization introduced in RFC patch D99750.

[vkmr]

Just a friendly ping!

[vkmr]

Rebasing on more recent upstream main.

[ABataev]

Tests?

llvm/include/llvm/IR/IRBuilder.h

2609

const

llvm/include/llvm/IR/IRBuilder.h

public:
 
   Value *CreatePreserveUnionAccessIndex(Value *Base, unsigned FieldIndex,
                                         MDNode *DbgInfo);
 
   Value *CreatePreserveStructAccessIndex(Type *ElTy, Value *Base,
                                          unsigned Index, unsigned FieldIndex,
                                          MDNode *DbgInfo);
 
+  /// Return an all true boolean vector of size and scalability \p NumElts.
+  Value *getTrueVector(ElementCount NumElts) {
+    VectorType *VTy = VectorType::get(Type::getInt1Ty(Context), NumElts);
+    return Constant::getAllOnesValue(VTy);
+  }
+
 private:
   /// Helper function that creates an assume intrinsic call that
   /// represents an alignment assumption on the provided pointer \p PtrValue
   /// with offset \p OffsetValue and alignment value \p AlignValue.
   CallInst *CreateAlignmentAssumptionHelper(const DataLayout &DL,
                                             Value *PtrValue, Value *AlignValue,
                                             Value *OffsetValue);
 

llvm/lib/Transforms/Vectorize/VPlan.h

public:
   /// Returns true if the recipe may read from or write to memory.
   bool mayReadOrWriteMemory() const {
     return mayReadFromMemory() || mayWriteToMemory();
   }
 };
 
 inline bool VPUser::classof(const VPDef *Def) {
   return Def->getVPDefID() == VPRecipeBase::VPInstructionSC ||
+         Def->getVPDefID() == VPRecipeBase::VPPredicatedWidenSC ||
          Def->getVPDefID() == VPRecipeBase::VPWidenSC ||
          Def->getVPDefID() == VPRecipeBase::VPWidenCallSC ||
          Def->getVPDefID() == VPRecipeBase::VPWidenSelectSC ||
          Def->getVPDefID() == VPRecipeBase::VPWidenGEPSC ||
          Def->getVPDefID() == VPRecipeBase::VPBlendSC ||
          Def->getVPDefID() == VPRecipeBase::VPInterleaveSC ||
          Def->getVPDefID() == VPRecipeBase::VPReplicateSC ||
          Def->getVPDefID() == VPRecipeBase::VPReductionSC ||
          Def->getVPDefID() == VPRecipeBase::VPBranchOnMaskSC ||
+         Def->getVPDefID() ==
+             VPRecipeBase::VPPredicatedWidenMemoryInstructionSC ||
          Def->getVPDefID() == VPRecipeBase::VPWidenMemoryInstructionSC;
 }
 
 /// This is a concrete Recipe that models a single VPlan-level instruction.
 /// While as any Recipe it may generate a sequence of IR instructions when
 /// executed, these instructions would always form a single-def expression as
 /// the VPInstruction is also a single def-use vertex.
 class VPInstruction : public VPRecipeBase, public VPValue {
   friend class VPlanSlp;
 
 public:
   /// VPlan opcodes, extending LLVM IR with idiomatics instructions.
   enum {
     FirstOrderRecurrenceSplice =
         Instruction::OtherOpsEnd + 1, // Combines the incoming and previous
                                       // values of a first-order recurrence.
     Not,
     ICmpULE,
     SLPLoad,
     SLPStore,
     ActiveLaneMask,
+    AllTrueMask,
   };
 
 private:
   typedef unsigned char OpcodeTy;
   OpcodeTy Opcode;
 
   /// Utility method serving execute(): generates a single instance of the
   /// modeled instruction.
bool hasResult() const {
     }
   }
 };
 
 /// VPWidenRecipe is a recipe for producing a copy of vector type its
 /// ingredient. This recipe covers most of the traditional vectorization cases
 /// where each ingredient transforms into a vectorized version of itself.
 class VPWidenRecipe : public VPRecipeBase, public VPValue {
+protected:
+  template <typename IterT>
+  VPWidenRecipe(Instruction &I, iterator_range<IterT> Operands,
+                const unsigned char RecipeSC, const unsigned char ValueSC)
+      : VPRecipeBase(RecipeSC, Operands), VPValue(ValueSC, &I, this) {}
+
 public:
   template <typename IterT>
   VPWidenRecipe(Instruction &I, iterator_range<IterT> Operands)
-      : VPRecipeBase(VPRecipeBase::VPWidenSC, Operands),
-        VPValue(VPValue::VPVWidenSC, &I, this) {}
+      : VPWidenRecipe(I, Operands, VPRecipeBase::VPWidenSC,
+                      VPValue::VPVWidenSC) {}
 
   ~VPWidenRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPWidenSC;
   }
   static inline bool classof(const VPValue *V) {
     return V->getVPValueID() == VPValue::VPVWidenSC;
   }
 
   /// Produce widened copies of all Ingredients.
   void execute(VPTransformState &State) override;
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
 };
 
+/// VPPredicatedWidenRecipe is a recipe for producing a copy of vector type
+/// using VP intrinsics for its ingredient. This recipe covers most of the
+/// traditional vectorization cases where each ingredient transforms into a
+/// vectorized version of itself.
+class VPPredicatedWidenRecipe : public VPWidenRecipe {
+public:
+  template <typename IterT>
+  VPPredicatedWidenRecipe(Instruction &I, iterator_range<IterT> Operands,
+                          VPValue *Mask, VPValue *EVL)
+      : VPWidenRecipe(I, Operands, VPRecipeBase::VPPredicatedWidenSC,
+                      VPValue::VPVPredicatedWidenSC) {
+    addOperand(Mask);
+    addOperand(EVL);
+  }
+
+  ~VPPredicatedWidenRecipe() override = default;
+
+  /// Method to support type inquiry through isa, cast, and dyn_cast.
+  static inline bool classof(const VPDef *D) {
+    return D->getVPDefID() == VPRecipeBase::VPPredicatedWidenSC;
+  }
+  static inline bool classof(const VPValue *V) {
+    return V->getVPValueID() == VPValue::VPVPredicatedWidenSC;
+  }
+
+  /// Return the mask used by this recipe.
+  VPValue *getMask() const { return getOperand(getNumOperands() - 2); }
+
+  /// Return the explicit vector length used by this recipe.
+  VPValue *getEVL() const { return getOperand(getNumOperands() - 1); }
+
+  /// Generate the wide load/store.
+  void execute(VPTransformState &State) override;
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// Print the recipe.
+  void print(raw_ostream &O, const Twine &Indent,
+             VPSlotTracker &SlotTracker) const override;
+#endif
+};
+
 /// A recipe for widening Call instructions.
 class VPWidenCallRecipe : public VPRecipeBase, public VPValue {
 
 public:
   template <typename IterT>
   VPWidenCallRecipe(CallInst &I, iterator_range<IterT> CallArguments)
       : VPRecipeBase(VPRecipeBase::VPWidenCallSC, CallArguments),
         VPValue(VPValue::VPVWidenCallSC, &I, this) {}
#endif
   /// by a nullptr.
   VPValue *getMask() const {
     assert(getNumOperands() <= 1 && "should have either 0 or 1 operands");
     // Mask is optional.
     return getNumOperands() == 1 ? getOperand(0) : nullptr;
   }
 };
 
+/// A recipe to generate Explicit Vector Length (EVL) value to be used with
+/// VPred intrinsics.
+class VPWidenEVLRecipe : public VPRecipeBase, public VPValue {
+
+public:
+  VPWidenEVLRecipe(VPValue *IV, VPValue *TC)
+      : VPRecipeBase(VPRecipeBase::VPWidenEVLSC, {IV, TC}),
+        VPValue(VPValue::VPVWidenEVLSC, nullptr, this) {}
+  ~VPWidenEVLRecipe() override = default;
+
+  /// Return the VPValue representing EVL.
+  const VPValue *getEVL() const { return this; }
+  VPValue *getEVL() { return this; }
+
+  /// Return VPValue representing Induction Variable.
+  VPValue *getIV() const { return getOperand(0); }
+
+  /// Return VPValue representing trip count.
+  VPValue *getTripCount() const { return getOperand(1); }
+
+  /// Method to support type inquiry through isa, cast, and dyn_cast.
+  static inline bool classof(const VPDef *D) {
+    return D->getVPDefID() == VPRecipeBase::VPWidenEVLSC;
+  }
+
+  /// Generate the instructions to compute EVL.
+  void execute(VPTransformState &State) override;
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// Print the recipe.
+  void print(raw_ostream &O, const Twine &Indent,
+             VPSlotTracker &SlotTracker) const override;
+#endif
+};
+
 /// VPPredInstPHIRecipe is a recipe for generating the phi nodes needed when
 /// control converges back from a Branch-on-Mask. The phi nodes are needed in
 /// order to merge values that are set under such a branch and feed their uses.
 /// The phi nodes can be scalar or vector depending on the users of the value.
 /// This recipe works in concert with VPBranchOnMaskRecipe.
 class VPPredInstPHIRecipe : public VPRecipeBase, public VPValue {
 public:
   /// Construct a VPPredInstPHIRecipe given \p PredInst whose value needs a phi

 
 /// A Recipe for widening load/store operations.
 /// The recipe uses the following VPValues:
 /// - For load: Address, optional mask
 /// - For store: Address, stored value, optional mask
 /// TODO: We currently execute only per-part unless a specific instance is
 /// provided.
 class VPWidenMemoryInstructionRecipe : public VPRecipeBase {
+  void setMask(VPValue *Mask) {
+    if (!Mask)
+      return;
+    addOperand(Mask);
+  }
+
+  bool isMasked() const {
+    return isStore() ? getNumOperands() == 3 : getNumOperands() == 2;
+  }
+
+protected:
   Instruction &Ingredient;
 
   // Whether the loaded-from / stored-to addresses are consecutive.
   bool Consecutive;
 
   // Whether the consecutive loaded/stored addresses are in reverse order.
   bool Reverse;
 
-  void setMask(VPValue *Mask) {
-    if (!Mask)
-      return;
-    addOperand(Mask);
+
+  VPWidenMemoryInstructionRecipe(LoadInst &Load, VPValue *Addr,
+                                 const unsigned char RecipeSC,
+                                 const unsigned char ValueSC, bool Consecutive,
+               bool Reverse)
+      : VPRecipeBase(RecipeSC, {Addr}), Ingredient(Load),
+  Consecutive(Consecutive), Reverse(Reverse) {
+    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
+    new VPValue(ValueSC, &Load, this);
   }
 
-  bool isMasked() const {
-    return isStore() ? getNumOperands() == 3 : getNumOperands() == 2;
+  VPWidenMemoryInstructionRecipe(StoreInst &Store, VPValue *Addr,
+                                 VPValue *StoredValue,
+                                 const unsigned char RecipeSC, bool Consecutive,
+                                 bool Reverse)
+      : VPRecipeBase(RecipeSC, {Addr, StoredValue}), Ingredient(Store),
+        Consecutive(Consecutive), Reverse(Reverse) {
+    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
   }
 
 public:
   VPWidenMemoryInstructionRecipe(LoadInst &Load, VPValue *Addr, VPValue *Mask,
                                  bool Consecutive, bool Reverse)
-      : VPRecipeBase(VPWidenMemoryInstructionSC, {Addr}), Ingredient(Load),
-        Consecutive(Consecutive), Reverse(Reverse) {
-    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
-    new VPValue(VPValue::VPVMemoryInstructionSC, &Load, this);
+      : VPWidenMemoryInstructionRecipe(Load, Addr, VPWidenMemoryInstructionSC,
+                                       VPValue::VPVMemoryInstructionSC,
+                                       Consecutive, Reverse) {
     setMask(Mask);
   }
 
   VPWidenMemoryInstructionRecipe(StoreInst &Store, VPValue *Addr,
                                  VPValue *StoredValue, VPValue *Mask,
                                  bool Consecutive, bool Reverse)
-      : VPRecipeBase(VPWidenMemoryInstructionSC, {Addr, StoredValue}),
-        Ingredient(Store), Consecutive(Consecutive), Reverse(Reverse) {
-    assert((Consecutive || !Reverse) && "Reverse implies consecutive");
+      : VPWidenMemoryInstructionRecipe(Store, Addr, StoredValue,
+                                       VPWidenMemoryInstructionSC, Consecutive,
+                                       Reverse) {
     setMask(Mask);
   }
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPWidenMemoryInstructionSC;
   }
 

 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
 };
 
+/// A Recipe for widening load/store operations to VP intrinsics.
+/// The recipe uses the following VPValues:
+/// - For load: Address, mask, EVL
+/// - For store: Address, stored value, mask, EVL
+class VPPredicatedWidenMemoryInstructionRecipe
+    : public VPWidenMemoryInstructionRecipe {
+
+public:
+  VPPredicatedWidenMemoryInstructionRecipe(LoadInst &Load, VPValue *Addr,
+                                           VPValue *Mask, VPValue *EVL,
+                                           bool Consecutive, bool Reverse)
+      : VPWidenMemoryInstructionRecipe(
+            Load, Addr, VPPredicatedWidenMemoryInstructionSC,
+            VPValue::VPVPredicatedMemoryInstructionSC, Consecutive, Reverse) {
+    addOperand(Mask);
+    addOperand(EVL);
+  }
+
+  VPPredicatedWidenMemoryInstructionRecipe(StoreInst &Store, VPValue *Addr,
+                                           VPValue *StoredValue, VPValue *Mask,
+                                           VPValue *EVL, bool Consecutive,
+                                           bool Reverse)
+      : VPWidenMemoryInstructionRecipe(Store, Addr, StoredValue,
+                                       VPPredicatedWidenMemoryInstructionSC,
+                                       Consecutive, Reverse) {
+    addOperand(Mask);
+    addOperand(EVL);
+  }
+
+  /// Method to support type inquiry through isa, cast, and dyn_cast.
+  static inline bool classof(const VPDef *D) {
+    return D->getVPDefID() ==
+           VPRecipeBase::VPPredicatedWidenMemoryInstructionSC;
+  }
+
+  /// Return the mask used by this recipe.
+  VPValue *getMask() const {
+    // Mask is the second last, mandatory operand.
+    return getOperand(getNumOperands() - 2);
+  }
+
+  /// Return the EVL used by this recipe.
+  VPValue *getEVL() const {
+    // EVL is the last, mandatory operand.
+    return getOperand(getNumOperands() - 1);
+  }
+
+  /// Generate the wide load/store.
+  void execute(VPTransformState &State) override;
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// Print the recipe.
+  void print(raw_ostream &O, const Twine &Indent,
+             VPSlotTracker &SlotTracker) const override;
+#endif
+};
+
 /// A Recipe for widening the canonical induction variable of the vector loop.
 class VPWidenCanonicalIVRecipe : public VPRecipeBase, public VPValue {
 public:
   VPWidenCanonicalIVRecipe()
       : VPRecipeBase(VPWidenCanonicalIVSC, {}),
         VPValue(VPValue::VPVWidenCanonicalIVSC, nullptr, this) {}
 
   ~VPWidenCanonicalIVRecipe() override = default;

llvm/lib/Transforms/Vectorize/VPlan.cpp

bool VPRecipeBase::mayHaveSideEffects() const {
   switch (getVPDefID()) {
   case VPBranchOnMaskSC:
     return false;
   case VPWidenIntOrFpInductionSC:
   case VPWidenCanonicalIVSC:
   case VPWidenPHISC:
   case VPBlendSC:
   case VPWidenSC:
+  case VPPredicatedWidenSC:
   case VPWidenGEPSC:
   case VPReductionSC:
   case VPWidenSelectSC: {
     const Instruction *I =
         dyn_cast_or_null<Instruction>(getVPSingleValue()->getUnderlyingValue());
     (void)I;
     assert((!I || !I->mayHaveSideEffects()) &&
            "underlying instruction has side-effects");
case VPInstruction::FirstOrderRecurrenceSplice: {
     if (!PartMinus1->getType()->isVectorTy()) {
       State.set(this, PartMinus1, Part);
     } else {
       Value *V2 = State.get(getOperand(1), Part);
       State.set(this, Builder.CreateVectorSplice(PartMinus1, V2, -1), Part);
     }
     break;
   }
+  case VPInstruction::AllTrueMask: {
+    Value *AllTrueMask = Builder.getTrueVector(State.VF);
+    State.set(this, AllTrueMask, Part);
+    break;
+  }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
   }
 }
 
 void VPInstruction::execute(VPTransformState &State) {
   assert(!State.Instance && "VPInstruction executing an Instance");
   for (unsigned Part = 0; Part < State.UF; ++Part)
case VPInstruction::SLPStore:
     O << "combined store";
     break;
   case VPInstruction::ActiveLaneMask:
     O << "active lane mask";
     break;
   case VPInstruction::FirstOrderRecurrenceSplice:
     O << "first-order splice";
     break;
+  case VPInstruction::AllTrueMask:
+    O << "all true mask";
+    break;
   default:
     O << Instruction::getOpcodeName(getOpcode());
   }
 
   for (const VPValue *Operand : operands()) {
     O << " ";
     Operand->printAsOperand(O, SlotTracker);
   }

 void VPWidenRecipe::print(raw_ostream &O, const Twine &Indent,
                           VPSlotTracker &SlotTracker) const {
   O << Indent << "WIDEN ";
   printAsOperand(O, SlotTracker);
   O << " = " << getUnderlyingInstr()->getOpcodeName() << " ";
   printOperands(O, SlotTracker);
 }
 
+void VPPredicatedWidenRecipe::print(raw_ostream &O, const Twine &Indent,
+                                    VPSlotTracker &SlotTracker) const {
+  O << Indent << "PREDICATED-WIDEN ";
+  printAsOperand(O, SlotTracker);
+  O << " = " << getUnderlyingInstr()->getOpcodeName() << " ";
+  printOperands(O, SlotTracker);
+}
+
 void VPWidenIntOrFpInductionRecipe::print(raw_ostream &O, const Twine &Indent,
                                           VPSlotTracker &SlotTracker) const {
   O << Indent << "WIDEN-INDUCTION";
   if (getTruncInst()) {
     O << "\\l\"";
     O << " +\n" << Indent << "\"  " << VPlanIngredient(IV) << "\\l\"";
     O << " +\n" << Indent << "\"  ";
     getVPValue(0)->printAsOperand(O, SlotTracker);
if (!isStore()) {
     O << " = ";
   }
   O << Instruction::getOpcodeName(Ingredient.getOpcode()) << " ";
 
   printOperands(O, SlotTracker);
 }
 #endif
 
+void VPPredicatedWidenMemoryInstructionRecipe::print(
+    raw_ostream &O, const Twine &Indent, VPSlotTracker &SlotTracker) const {
+  O << Indent << "PREDICATED-WIDEN ";
+
+  if (!isStore()) {
+    getVPSingleValue()->printAsOperand(O, SlotTracker);
+    O << " = ";
+  }
+  O << Instruction::getOpcodeName(Ingredient.getOpcode()) << " ";
+
+  printOperands(O, SlotTracker);
+}
+
 void VPWidenCanonicalIVRecipe::execute(VPTransformState &State) {
   Value *CanonicalIV = State.CanonicalIV;
   Type *STy = CanonicalIV->getType();
   IRBuilder<> Builder(State.CFG.PrevBB->getTerminator());
   ElementCount VF = State.VF;
   assert(!VF.isScalable() && "the code following assumes non scalables ECs");
   Value *VStart = VF.isScalar()
                       ? CanonicalIV
void VPReductionPHIRecipe::print(raw_ostream &O, const Twine &Indent,
   O << Indent << "WIDEN-REDUCTION-PHI ";
 
   printAsOperand(O, SlotTracker);
   O << " = phi ";
   printOperands(O, SlotTracker);
 }
 #endif
 
+void VPWidenEVLRecipe::print(raw_ostream &O, const Twine &Indent,
+                             VPSlotTracker &SlotTracker) const {
+  O << Indent << "EMIT ";
+  getEVL()->printAsOperand(O, SlotTracker);
+  O << " = GENERATE-EXPLICIT-VECTOR-LENGTH";
+}
+
 template void DomTreeBuilder::Calculate<VPDominatorTree>(VPDominatorTree &DT);
 
 void VPValue::replaceAllUsesWith(VPValue *New) {
   for (unsigned J = 0; J < getNumUsers();) {
     VPUser *User = Users[J];
     unsigned NumUsers = getNumUsers();
     for (unsigned I = 0, E = User->getNumOperands(); I < E; ++I)
       if (User->getOperand(I) == this)

llvm/lib/Transforms/Vectorize/VPlanValue.h

class VPValue {
   friend class VPDef;
   friend class VPInstruction;
   friend struct VPlanTransforms;
   friend class VPBasicBlock;
   friend class VPInterleavedAccessInfo;
   friend class VPSlotTracker;
   friend class VPRecipeBase;
   friend class VPWidenMemoryInstructionRecipe;
+  friend class VPPredicatedWidenMemoryInstructionRecipe;
 
   const unsigned char SubclassID; ///< Subclass identifier (for isa/dyn_cast).
 
   SmallVector<VPUser *, 1> Users;
 
 protected:
   // Hold the underlying Value, if any, attached to this VPValue.
   Value *UnderlyingVal;
public:
   /// An enumeration for keeping track of the concrete subclass of VPValue that
   /// are actually instantiated. Values of this enumeration are kept in the
   /// SubclassID field of the VPValue objects. They are used for concrete
   /// type identification.
   enum {
     VPValueSC,
     VPVInstructionSC,
     VPVMemoryInstructionSC,
+    VPVPredicatedMemoryInstructionSC,
+    VPVPredicatedWidenSC,
     VPVReductionSC,
     VPVReplicateSC,
     VPVWidenSC,
     VPVWidenCallSC,
+    VPVWidenEVLSC,
     VPVWidenGEPSC,
     VPVWidenSelectSC,
 
     // Phi-like VPValues. Need to be kept together.
     VPVBlendSC,
     VPVFirstOrderRecurrencePHISC,
     VPVWidenPHISC,
     VPVWidenCanonicalIVSC,
public:
   /// An enumeration for keeping track of the concrete subclass of VPRecipeBase
   /// that is actually instantiated. Values of this enumeration are kept in the
   /// SubclassID field of the VPRecipeBase objects. They are used for concrete
   /// type identification.
   using VPRecipeTy = enum {
     VPBranchOnMaskSC,
     VPInstructionSC,
     VPInterleaveSC,
+    VPPredicatedWidenMemoryInstructionSC,
+    VPPredicatedWidenSC,
     VPReductionSC,
     VPReplicateSC,
     VPWidenCallSC,
+    VPWidenEVLSC,
     VPWidenGEPSC,
     VPWidenMemoryInstructionSC,
     VPWidenSC,
     VPWidenSelectSC,
 
     // Phi-like recipes. Need to be kept together.
     VPBlendSC,
     VPFirstOrderRecurrencePHISC,