[VPlan] Manage compare predicates in VPRecipeWithIRFlags.

[fhahn]

Summary

Extend VPRecipeWithIRFlags to also manage predicates for compares. This
allows removing the custom ICmpULE opcode from VPInstruction which was a
workaround for missing proper predicate handling.

This simplifies the code a bit while also allowing compares with any
predicates. It also fixes a case where the compare predixcate wasn't
printed properly for VPReplicateRecipes.

Discussed/split off from D150398.

[Ayal]

Thanks for following up on this! Adding various nits.

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

122–123

Add documentation? TODO add createFCmp() when needed?

Perhaps better place below, close to createNot/And/Or/Select, after the more general createNaryOp().
Have it returned createInstruction() like them, with a new createInstruction() that calls the new VPInstruction constructor?

BTW, the createNaryOp() with Inst seems a bit odd?

llvm/lib/Transforms/Vectorize/VPlan.h

817

nit: place after FPMathOp, to be consistent with the union order below?

860

nit: keep AllFlags last? An unsigned char still suffices, right?

877

nit: place below after handling FPMathOperator, for consistency.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

120

Is FCmp case dead, or checked by any test?

OTOH, would be ideal to generalize by checking if Instruction may have side effects for all opcodes up to Instruction::OtherOpsEnd, but w/o having an Instruction instance...

629

nit: place below for consistency.

llvm/test/Transforms/LoopVectorize/vplan-sink-scalars-and-merge.ll

909

nice catch!

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.h

1075–1076

Also have clone() replicate the IR flags? (Independent of this patch, calls for testing?)

[tschuett]

llvm/lib/Transforms/Vectorize/VPlan.h

854–855

No std::variant?

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

122–123

Added TODO and documentation and moved.

Using createInstruction would require adding a new variant of that function I think. Maybe it would be better to just add a helper to insert the created instruction.

BTW, the createNaryOp() with Inst seems a bit odd?

Used for VPlan-native path I think. Should be able to remove it once we add an in-place IR recipe.

llvm/lib/Transforms/Vectorize/VPlan.h

817

Moved, thanks!

854–855

Hmm, can be changed in the future. It looks like there are very few uses of std::variant in the llvm subdirectory so it doesn't seem used a lot.

860

No, updated the unsigned thanks! (Although all predicates should fit into a char...)

877

Updated to have Cmp at top of the enum and union.

1075–1076

Seemed unused, will remove separately.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

120

Dead code for now, removed!

OTOH, would be ideal to generalize by checking if Instruction may have side effects for all opcodes up to Instruction::OtherOpsEnd, but w/o having an Instruction instance...

Agreed, that would require refactoring the corresponding implementation in Instruction to just work on opcodes.

629

Reordered code to have Cmp at front everywhere.

[Ayal]

Reordered code to have Cmp at front everywhere.

Fine, let's have it at front everywhere.

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

49

nit: tryInsertInstruction?

llvm/lib/Transforms/Vectorize/VPlan.h

910

nit: place this first as well

999

nit: let Cmp appear first consistently

1074

nit: ditto

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

120

Perhaps worth leaving behind a TODO.
Perhaps native may find this useful already?

267

nit: ditto

283

nit: FCmp is currently dead?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

49

Adjusted, thanks!

llvm/lib/Transforms/Vectorize/VPlan.h

910

Moved, thanks!

999

Moved, thanks!

1074

Moved, thanks!

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

267

Moved, thanks!

283

dropped, thanks!

[fhahn]

Address comments, thanks!

[Ayal]

LGTM, adding last couple of nits. thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

163

nit: assert Pred is an Integer compare predicate.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

255

nit: assert Opcode is an ICmp; introduce support for FCmp later when needed.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

163

Added in the committed version, thanks!

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

255

Removed in the committed version, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

 class TargetLibraryInfo;
 class VPRecipeBuilder;
 
 /// VPlan-based builder utility analogous to IRBuilder.
 class VPBuilder {
   VPBasicBlock *BB = nullptr;
   VPBasicBlock::iterator InsertPt = VPBasicBlock::iterator();
 
+  /// Insert \p VPI in BB at InsertPt if BB is set.
+  VPInstruction *tryInsertInstruction(VPInstruction *VPI) {
+    if (BB)
+      BB->insert(VPI, InsertPt);
+    return VPI;
+  }
+
   VPInstruction *createInstruction(unsigned Opcode,
                                    ArrayRef<VPValue *> Operands, DebugLoc DL,
                                    const Twine &Name = "") {
-    VPInstruction *Instr = new VPInstruction(Opcode, Operands, DL, Name);
-    if (BB)
-      BB->insert(Instr, InsertPt);
-    return Instr;
+    return tryInsertInstruction(new VPInstruction(Opcode, Operands, DL, Name));
   }
 
   VPInstruction *createInstruction(unsigned Opcode,
                                    std::initializer_list<VPValue *> Operands,
                                    DebugLoc DL, const Twine &Name = "") {
     return createInstruction(Opcode, ArrayRef<VPValue *>(Operands), DL, Name);
   }
 
public:
   }
 
   /// This specifies that created instructions should be inserted at the
   /// specified point.
   void setInsertPoint(VPBasicBlock *TheBB, VPBasicBlock::iterator IP) {
     BB = TheBB;
     InsertPt = IP;
   }
 
   /// Create an N-ary operation with \p Opcode, \p Operands and set \p Inst as
   /// its underlying Instruction.
   VPValue *createNaryOp(unsigned Opcode, ArrayRef<VPValue *> Operands,
                         Instruction *Inst = nullptr, const Twine &Name = "") {
     DebugLoc DL;
     if (Inst)
       DL = Inst->getDebugLoc();
     VPInstruction *NewVPInst = createInstruction(Opcode, Operands, DL, Name);
     NewVPInst->setUnderlyingValue(Inst);
public:
   }
 
   VPValue *createSelect(VPValue *Cond, VPValue *TrueVal, VPValue *FalseVal,
                         DebugLoc DL, const Twine &Name = "") {
     return createNaryOp(Instruction::Select, {Cond, TrueVal, FalseVal}, DL,
                         Name);
   }
 
+  /// Create a new ICmp VPInstruction with predicate \p Pred and operands \p A
+  /// and \p B.
+  /// TODO: add createFCmp when needed.
+  VPValue *createICmp(CmpInst::Predicate Pred, VPValue *A, VPValue *B,
+                      DebugLoc DL = {}, const Twine &Name = "") {
+    return tryInsertInstruction(
+        new VPInstruction(Instruction::ICmp, Pred, A, B, DL, Name));
+  }
+
   //===--------------------------------------------------------------------===//
   // RAII helpers.
   //===--------------------------------------------------------------------===//
 
   /// RAII object that stores the current insertion point and restores it when
   /// the object is destroyed.
   class InsertPointGuard {
     VPBuilder &Builder;

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void VPRecipeBuilder::createHeaderMask(VPlan &Plan) {
   Builder.setInsertPoint(HeaderVPBB, NewInsertionPoint);
   VPValue *BlockMask = nullptr;
   if (useActiveLaneMask(TFStyle)) {
     VPValue *TC = Plan.getTripCount();
     BlockMask = Builder.createNaryOp(VPInstruction::ActiveLaneMask, {IV, TC},
                                      nullptr, "active.lane.mask");
   } else {
     VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
-    BlockMask = Builder.createNaryOp(VPInstruction::ICmpULE, {IV, BTC});
+    BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
   }
   BlockMaskCache[Header] = BlockMask;
 }
 
 VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlan &Plan) {
   assert(OrigLoop->contains(BB) && "Block is not a part of a loop");
 
   // Look for cached value.

llvm/lib/Transforms/Vectorize/VPlan.h

#define VP_CLASSOF_IMPL(VPDefID)                                               \
   }                                                                            \
   static inline bool classof(const VPRecipeBase *R) {                          \
     return R->getVPDefID() == VPDefID;                                         \
   }
 
 /// Class to record LLVM IR flag for a recipe along with it.
 class VPRecipeWithIRFlags : public VPRecipeBase {
   enum class OperationType : unsigned char {
+    Cmp,
     OverflowingBinOp,
     PossiblyExactOp,
     GEPOp,
     FPMathOp,
     Other
   };
 
 public:
struct FastMathFlagsTy {
     char AllowContract : 1;
     char ApproxFunc : 1;
 
     FastMathFlagsTy(const FastMathFlags &FMF);
   };
 
   OperationType OpType;
 
   union {
+    CmpInst::Predicate CmpPredicate;
     WrapFlagsTy WrapFlags;
     ExactFlagsTy ExactFlags;
     GEPFlagsTy GEPFlags;
     FastMathFlagsTy FMFs;
-    unsigned char AllFlags;
+    unsigned AllFlags;
   };
 
 public:
   template <typename IterT>
   VPRecipeWithIRFlags(const unsigned char SC, IterT Operands)
       : VPRecipeBase(SC, Operands) {
     OpType = OperationType::Other;
     AllFlags = 0;
   }
 
   template <typename IterT>
   VPRecipeWithIRFlags(const unsigned char SC, IterT Operands, Instruction &I)
       : VPRecipeWithIRFlags(SC, Operands) {
-    if (auto *Op = dyn_cast<OverflowingBinaryOperator>(&I)) {
+    if (auto *Op = dyn_cast<CmpInst>(&I)) {
+      OpType = OperationType::Cmp;
+      CmpPredicate = Op->getPredicate();
+    } else if (auto *Op = dyn_cast<OverflowingBinaryOperator>(&I)) {
       OpType = OperationType::OverflowingBinOp;
       WrapFlags = {Op->hasNoUnsignedWrap(), Op->hasNoSignedWrap()};
     } else if (auto *Op = dyn_cast<PossiblyExactOperator>(&I)) {
       OpType = OperationType::PossiblyExactOp;
       ExactFlags.IsExact = Op->isExact();
     } else if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
       OpType = OperationType::GEPOp;
       GEPFlags.IsInBounds = GEP->isInBounds();
     } else if (auto *Op = dyn_cast<FPMathOperator>(&I)) {
       OpType = OperationType::FPMathOp;
       FMFs = Op->getFastMathFlags();
     }
   }
 
   template <typename IterT>
   VPRecipeWithIRFlags(const unsigned char SC, IterT Operands,
+                      CmpInst::Predicate Pred)
+      : VPRecipeBase(SC, Operands), OpType(OperationType::Cmp),
+        CmpPredicate(Pred) {}
+
+  template <typename IterT>
+  VPRecipeWithIRFlags(const unsigned char SC, IterT Operands,
                       WrapFlagsTy WrapFlags)
       : VPRecipeBase(SC, Operands), OpType(OperationType::OverflowingBinOp),
         WrapFlags(WrapFlags) {}
 
   template <typename IterT>
   VPRecipeWithIRFlags(const unsigned char SC, IterT Operands,
                       FastMathFlags FMFs)
       : VPRecipeBase(SC, Operands), OpType(OperationType::FPMathOp),
         FMFs(FMFs) {}
 
   static inline bool classof(const VPRecipeBase *R) {
     return R->getVPDefID() == VPRecipeBase::VPInstructionSC ||
            R->getVPDefID() == VPRecipeBase::VPWidenSC ||
            R->getVPDefID() == VPRecipeBase::VPWidenGEPSC ||
            R->getVPDefID() == VPRecipeBase::VPReplicateSC;
   }
 
   /// Drop all poison-generating flags.
   void dropPoisonGeneratingFlags() {
case OperationType::PossiblyExactOp:
       break;
     case OperationType::GEPOp:
       GEPFlags.IsInBounds = false;
       break;
     case OperationType::FPMathOp:
       FMFs.NoNaNs = false;
       FMFs.NoInfs = false;
       break;
+    case OperationType::Cmp:
     case OperationType::Other:
       break;
     }
   }
 
   /// Set the IR flags for \p I.
   void setFlags(Instruction *I) const {
     switch (OpType) {
case OperationType::FPMathOp:
       I->setHasAllowReassoc(FMFs.AllowReassoc);
       I->setHasNoNaNs(FMFs.NoNaNs);
       I->setHasNoInfs(FMFs.NoInfs);
       I->setHasNoSignedZeros(FMFs.NoSignedZeros);
       I->setHasAllowReciprocal(FMFs.AllowReciprocal);
       I->setHasAllowContract(FMFs.AllowContract);
       I->setHasApproxFunc(FMFs.ApproxFunc);
       break;
+    case OperationType::Cmp:
     case OperationType::Other:
       break;
     }
   }
 
+  CmpInst::Predicate getPredicate() const {
+    assert(OpType == OperationType::Cmp &&
+           "recipe doesn't have a compare predicate");
+    return CmpPredicate;
+  }
+
   bool isInBounds() const {
     assert(OpType == OperationType::GEPOp &&
            "recipe doesn't have inbounds flag");
     return GEPFlags.IsInBounds;
   }
 
   /// Returns true if the recipe has fast-math flags.
   bool hasFastMathFlags() const { return OpType == OperationType::FPMathOp; }
 
   FastMathFlags getFastMathFlags() const;
 
   bool hasNoUnsignedWrap() const {
     assert(OpType == OperationType::OverflowingBinOp &&
            "recipe doesn't have a NUW flag");
     return WrapFlags.HasNUW;
   }
 
   bool hasNoSignedWrap() const {
     assert(OpType == OperationType::OverflowingBinOp &&
            "recipe doesn't have a NSW flag");
     return WrapFlags.HasNSW;
   }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   void printFlags(raw_ostream &O) const;
 #endif
 };
 
 /// 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 VPRecipeWithIRFlags, 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,
     CalculateTripCountMinusVF,
     CanonicalIVIncrement,
     // The next op is similar to the above, but instead increment the
     // canonical IV separately for each unrolled part.
     CanonicalIVIncrementForPart,
VPInstruction(unsigned Opcode, ArrayRef<VPValue *> Operands, DebugLoc DL,
                 const Twine &Name = "")
       : VPRecipeWithIRFlags(VPDef::VPInstructionSC, Operands), VPValue(this),
         Opcode(Opcode), DL(DL), Name(Name.str()) {}
 
   VPInstruction(unsigned Opcode, std::initializer_list<VPValue *> Operands,
                 DebugLoc DL = {}, const Twine &Name = "")
       : VPInstruction(Opcode, ArrayRef<VPValue *>(Operands), DL, Name) {}
 
+  VPInstruction(unsigned Opcode, CmpInst::Predicate Pred, VPValue *A,
+                VPValue *B, DebugLoc DL = {}, const Twine &Name = "");
+
   VPInstruction(unsigned Opcode, std::initializer_list<VPValue *> Operands,
                 WrapFlagsTy WrapFlags, DebugLoc DL = {}, const Twine &Name = "")
       : VPRecipeWithIRFlags(VPDef::VPInstructionSC, Operands, WrapFlags),
         VPValue(this), Opcode(Opcode), DL(DL), Name(Name.str()) {}
 
   VPInstruction(unsigned Opcode, std::initializer_list<VPValue *> Operands,
                 FastMathFlags FMFs, DebugLoc DL = {}, const Twine &Name = "");
 
   VP_CLASSOF_IMPL(VPDef::VPInstructionSC)
 
   unsigned getOpcode() const { return Opcode; }
 
   /// Generate the instruction.
   /// TODO: We currently execute only per-part unless a specific instance is
   /// provided.
   void execute(VPTransformState &State) override;
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the VPInstruction to \p O.

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

 
 bool VPRecipeBase::mayHaveSideEffects() const {
   switch (getVPDefID()) {
   case VPDerivedIVSC:
   case VPPredInstPHISC:
     return false;
   case VPInstructionSC:
     switch (cast<VPInstruction>(this)->getOpcode()) {
+    case Instruction::ICmp:
     case VPInstruction::Not:
-    case VPInstruction::ICmpULE:
     case VPInstruction::CalculateTripCountMinusVF:
     case VPInstruction::CanonicalIVIncrement:
     case VPInstruction::CanonicalIVIncrementForPart:
       return false;
     default:
       return true;
     }
   case VPWidenCallSC:
FastMathFlags VPRecipeWithIRFlags::getFastMathFlags() const {
   Res.setNoInfs(FMFs.NoInfs);
   Res.setNoSignedZeros(FMFs.NoSignedZeros);
   Res.setAllowReciprocal(FMFs.AllowReciprocal);
   Res.setAllowContract(FMFs.AllowContract);
   Res.setApproxFunc(FMFs.ApproxFunc);
   return Res;
 }
 
+VPInstruction::VPInstruction(unsigned Opcode, CmpInst::Predicate Pred,
+                             VPValue *A, VPValue *B, DebugLoc DL,
+                             const Twine &Name)
+    : VPRecipeWithIRFlags(VPDef::VPInstructionSC, ArrayRef<VPValue *>({A, B}),
+                          Pred),
+      VPValue(this), Opcode(Opcode), DL(DL), Name(Name.str()) {
+  assert(Opcode == Instruction::ICmp || Opcode == Instruction::FCmp);
+}
+
 VPInstruction::VPInstruction(unsigned Opcode,
                              std::initializer_list<VPValue *> Operands,
                              FastMathFlags FMFs, DebugLoc DL, const Twine &Name)
     : VPRecipeWithIRFlags(VPDef::VPInstructionSC, Operands, FMFs),
       VPValue(this), Opcode(Opcode), DL(DL), Name(Name.str()) {
   // Make sure the VPInstruction is a floating-point operation.
   assert(isFPMathOp() && "this op can't take fast-math flags");
 }
 
 Value *VPInstruction::generateInstruction(VPTransformState &State,
                                           unsigned Part) {
   IRBuilderBase &Builder = State.Builder;
   Builder.SetCurrentDebugLocation(DL);
 
   if (Instruction::isBinaryOp(getOpcode())) {
     Value *A = State.get(getOperand(0), Part);
     Value *B = State.get(getOperand(1), Part);
     return Builder.CreateBinOp((Instruction::BinaryOps)getOpcode(), A, B, Name);
   }
 
   switch (getOpcode()) {
   case VPInstruction::Not: {
     Value *A = State.get(getOperand(0), Part);
     return Builder.CreateNot(A, Name);
   }
-  case VPInstruction::ICmpULE: {
+  case Instruction::ICmp: {
     Value *A = State.get(getOperand(0), Part);
     Value *B = State.get(getOperand(1), Part);
-    return Builder.CreateICmpULE(A, B, Name);
+    return Builder.CreateCmp(getPredicate(), A, B, Name);
   }
   case Instruction::Select: {
     Value *Cond = State.get(getOperand(0), Part);
     Value *Op1 = State.get(getOperand(1), Part);
     Value *Op2 = State.get(getOperand(2), Part);
     return Builder.CreateSelect(Cond, Op1, Op2, Name);
   }
   case VPInstruction::ActiveLaneMask: {
if (hasResult()) {
     printAsOperand(O, SlotTracker);
     O << " = ";
   }
 
   switch (getOpcode()) {
   case VPInstruction::Not:
     O << "not";
     break;
-  case VPInstruction::ICmpULE:
-    O << "icmp ule";
-    break;
   case VPInstruction::SLPLoad:
     O << "combined load";
     break;
   case VPInstruction::SLPStore:
     O << "combined store";
     break;
   case VPInstruction::ActiveLaneMask:
     O << "active lane mask";
VPRecipeWithIRFlags::FastMathFlagsTy::FastMathFlagsTy(
   AllowReciprocal = FMF.allowReciprocal();
   AllowContract = FMF.allowContract();
   ApproxFunc = FMF.approxFunc();
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPRecipeWithIRFlags::printFlags(raw_ostream &O) const {
   switch (OpType) {
+  case OperationType::Cmp:
+    O << " " << CmpInst::getPredicateName(getPredicate());
+    break;
   case OperationType::PossiblyExactOp:
     if (ExactFlags.IsExact)
       O << " exact";
     break;
   case OperationType::OverflowingBinOp:
     if (WrapFlags.HasNUW)
       O << " nuw";
     if (WrapFlags.HasNSW)

 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPWidenRecipe::print(raw_ostream &O, const Twine &Indent,
                           VPSlotTracker &SlotTracker) const {
   O << Indent << "WIDEN ";
   printAsOperand(O, SlotTracker);
   const Instruction *UI = getUnderlyingInstr();
   O << " = " << UI->getOpcodeName();
   printFlags(O);
-  if (auto *Cmp = dyn_cast<CmpInst>(UI))
-    O << Cmp->getPredicate() << " ";
   printOperands(O, SlotTracker);
 }
 #endif
 
 void VPWidenCastRecipe::execute(VPTransformState &State) {
   auto *I = cast_or_null<Instruction>(getUnderlyingValue());
   if (I)
     State.setDebugLocFromInst(I);

llvm/test/Transforms/LoopVectorize/vplan-sink-scalars-and-merge.ll

 ; CHECK-NEXT: <xVFxUF> pred.store: {
 ; CHECK-NEXT:   pred.store.entry:
 ; CHECK-NEXT:     BRANCH-ON-MASK ir<%c>
 ; CHECK-NEXT:   Successor(s): pred.store.if, pred.store.continue
 ; CHECK-EMPTY:
 ; CHECK-NEXT:   pred.store.if:
 ; CHECK-NEXT:     REPLICATE ir<%l1> = load ir<%src>
 ; CHECK-NEXT:     REPLICATE ir<%l2> = trunc ir<%l1>
-; CHECK-NEXT:     REPLICATE ir<%cmp> = icmp ir<%l1>, ir<0>
+; CHECK-NEXT:     REPLICATE ir<%cmp> = icmp eq ir<%l1>, ir<0>
 ; CHECK-NEXT:     REPLICATE ir<%sel> = select ir<%cmp>, ir<5>, ir<%l2>
 ; CHECK-NEXT:     REPLICATE store ir<%sel>, ir<%dst>
 ; CHECK-NEXT:   Successor(s): pred.store.continue
 ; CHECK-EMPTY:
 ; CHECK-NEXT:   pred.store.continue:
 ; CHECK-NEXT:     PHI-PREDICATED-INSTRUCTION vp<[[PRED:%.+]]> = ir<%l1>
 ; CHECK-NEXT:   No successors
 ; CHECK-NEXT: }