VPlan: use Worklist in simplifyRecipes

Since simplifyRecipe creates new recipes in some cases, use a Worklist
in its caller to capture newly-created recipes, and add it to the
Worklist, as a candidate for further simplification. This patch
thoroughly rewrites simplifyRecipe to simplify matched patterns,
eraseFromParent when applicable, and simplify the logic.

Author: artagnon

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

@@ -161,33 +161,34 @@ class VPBuilder {
     return tryInsertInstruction(
         new VPInstruction(Opcode, Operands, WrapFlags, DL, Name));
   }
-  VPValue *createNot(VPValue *Operand, DebugLoc DL = {},
-                     const Twine &Name = "") {
+  VPInstruction *createNot(VPValue *Operand, DebugLoc DL = {},
+                           const Twine &Name = "") {
     return createInstruction(VPInstruction::Not, {Operand}, DL, Name);
   }
 
-  VPValue *createAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
-                     const Twine &Name = "") {
+  VPInstruction *createAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
+                           const Twine &Name = "") {
     return createInstruction(Instruction::BinaryOps::And, {LHS, RHS}, DL, Name);
   }
 
-  VPValue *createOr(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
-                    const Twine &Name = "") {
+  VPInstruction *createOr(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
+                          const Twine &Name = "") {
 
     return tryInsertInstruction(new VPInstruction(
         Instruction::BinaryOps::Or, {LHS, RHS},
         VPRecipeWithIRFlags::DisjointFlagsTy(false), DL, Name));
   }
 
-  VPValue *createLogicalAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
-                            const Twine &Name = "") {
+  VPInstruction *createLogicalAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
+                                  const Twine &Name = "") {
     return tryInsertInstruction(
         new VPInstruction(VPInstruction::LogicalAnd, {LHS, RHS}, DL, Name));
   }
 
-  VPValue *createSelect(VPValue *Cond, VPValue *TrueVal, VPValue *FalseVal,
-                        DebugLoc DL = {}, const Twine &Name = "",
-                        std::optional<FastMathFlags> FMFs = std::nullopt) {
+  VPInstruction *
+  createSelect(VPValue *Cond, VPValue *TrueVal, VPValue *FalseVal,
+               DebugLoc DL = {}, const Twine &Name = "",
+               std::optional<FastMathFlags> FMFs = std::nullopt) {
     auto *Select =
         FMFs ? new VPInstruction(Instruction::Select, {Cond, TrueVal, FalseVal},
                                  *FMFs, DL, Name)
@@ -199,8 +200,8 @@ class VPBuilder {
   /// 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 = "");
+  VPInstruction *createICmp(CmpInst::Predicate Pred, VPValue *A, VPValue *B,
+                            DebugLoc DL = {}, const Twine &Name = "");
 
   //===--------------------------------------------------------------------===//
   // RAII helpers.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -6903,8 +6903,9 @@ void LoopVectorizationCostModel::collectInLoopReductions() {
   }
 }
 
-VPValue *VPBuilder::createICmp(CmpInst::Predicate Pred, VPValue *A, VPValue *B,
-                               DebugLoc DL, const Twine &Name) {
+VPInstruction *VPBuilder::createICmp(CmpInst::Predicate Pred, VPValue *A,
+                                     VPValue *B, DebugLoc DL,
+                                     const Twine &Name) {
   assert(Pred >= CmpInst::FIRST_ICMP_PREDICATE &&
          Pred <= CmpInst::LAST_ICMP_PREDICATE && "invalid predicate");
   return tryInsertInstruction(

llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h

@@ -75,8 +75,9 @@ template <unsigned BitWidth = 0> struct specific_intval {
     if (!CI)
       return false;
 
-    assert((BitWidth == 0 || CI->getBitWidth() == BitWidth) &&
-           "Trying the match constant with unexpected bitwidth.");
+    if (BitWidth != 0 && CI->getBitWidth() != BitWidth)
+      return false;
+
     return APInt::isSameValue(CI->getValue(), Val);
   }
 };
@@ -87,6 +88,8 @@ inline specific_intval<0> m_SpecificInt(uint64_t V) {
 
 inline specific_intval<1> m_False() { return specific_intval<1>(APInt(64, 0)); }
 
+inline specific_intval<1> m_True() { return specific_intval<1>(APInt(64, 1)); }
+
 /// Matching combinators
 template <typename LTy, typename RTy> struct match_combine_or {
   LTy L;

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -24,6 +24,7 @@
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PatternMatch.h"
+#include <deque>
 
 using namespace llvm;
 
@@ -892,9 +893,10 @@ void VPlanTransforms::clearReductionWrapFlags(VPlan &Plan) {
   }
 }
 
-/// Try to simplify recipe \p R. Returns any new recipes introduced during
-/// simplification, as a candidate for further simplification.
-static VPRecipeBase *simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
+/// Try to simplify recipe \p R.  Returns any new recipes introduced during
+/// simplification, as candidates for further simplification.
+static SmallVector<VPRecipeBase *>
+simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo, VPlan &Plan) {
   using namespace llvm::VPlanPatternMatch;
 
   if (auto *Blend = dyn_cast<VPBlendRecipe>(&R)) {
@@ -909,11 +911,11 @@ static VPRecipeBase *simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
     if (UniqueValues.size() == 1) {
       Blend->replaceAllUsesWith(*UniqueValues.begin());
       Blend->eraseFromParent();
-      return nullptr;
+      return {};
     }
 
     if (Blend->isNormalized())
-      return nullptr;
+      return {};
 
     // Normalize the blend so its first incomming value is used as the initial
     // value with the others blended into it.
@@ -937,7 +939,7 @@ static VPRecipeBase *simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
     Blend->replaceAllUsesWith(NewBlend);
     Blend->eraseFromParent();
     recursivelyDeleteDeadRecipes(DeadMask);
-    return nullptr;
+    return {};
   }
 
   VPValue *A;
@@ -986,25 +988,74 @@ static VPRecipeBase *simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
         assert(TypeInfo.inferScalarType(VPV) == TypeInfo2.inferScalarType(VPV));
     }
 #endif
-    return VPC;
+    if (VPC)
+      return {VPC};
+    return {};
+  }
+
+  VPValue *X, *X1, *Y, *Z;
+  LLVMContext &Ctx = TypeInfo.getContext();
+
+  // (X || !X) -> true.
+  if (match(&R, m_c_BinaryOr(m_VPValue(X), m_Not(m_VPValue(X1)))) && X == X1) {
+    VPValue *VPV = Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx));
+    R.getVPSingleValue()->replaceAllUsesWith(VPV);
+    return {};
+  }
+
+  // (X || true) -> true.
+  if (match(&R, m_c_BinaryOr(m_VPValue(X), m_True()))) {
+    VPValue *VPV = Plan.getOrAddLiveIn(ConstantInt::getTrue(Ctx));
+    R.getVPSingleValue()->replaceAllUsesWith(VPV);
+    return {};
   }
 
-  // Simplify (X && Y) || (X && !Y) -> X.
-  // TODO: Split up into simpler, modular combines: (X && Y) || (X && Z) into X
-  // && (Y || Z) and (X || !X) into true. This requires queuing newly created
-  // recipes to be visited during simplification.
-  VPValue *X, *Y, *X1, *Y1;
-  if (match(&R,
-            m_c_BinaryOr(m_LogicalAnd(m_VPValue(X), m_VPValue(Y)),
-                         m_LogicalAnd(m_VPValue(X1), m_Not(m_VPValue(Y1))))) &&
-      X == X1 && Y == Y1) {
+  // (X || false) -> X.
+  if (match(&R, m_c_BinaryOr(m_VPValue(X), m_False()))) {
     R.getVPSingleValue()->replaceAllUsesWith(X);
-    return nullptr;
+    return {};
   }
 
-  if (match(&R, m_c_Mul(m_VPValue(A), m_SpecificInt(1))))
-    R.getVPSingleValue()->replaceAllUsesWith(A);
-  return nullptr;
+  // (X && !X) -> false.
+  if (match(&R, m_LogicalAnd(m_VPValue(X), m_Not(m_VPValue(X1)))) && X == X1) {
+    VPValue *VPV = Plan.getOrAddLiveIn(ConstantInt::getFalse(Ctx));
+    R.getVPSingleValue()->replaceAllUsesWith(VPV);
+    return {};
+  }
+
+  // (X && true) -> X.
+  if (match(&R, m_LogicalAnd(m_VPValue(X), m_True()))) {
+    R.getVPSingleValue()->replaceAllUsesWith(X);
+    return {};
+  }
+
+  // (X && false) -> false.
+  if (match(&R, m_LogicalAnd(m_VPValue(X), m_False()))) {
+    VPValue *VPV = Plan.getOrAddLiveIn(ConstantInt::getFalse(Ctx));
+    R.getVPSingleValue()->replaceAllUsesWith(VPV);
+    return {};
+  }
+
+  // (X * 1) -> X.
+  if (match(&R, m_c_Mul(m_VPValue(X), m_SpecificInt(1)))) {
+    R.getVPSingleValue()->replaceAllUsesWith(X);
+    return {};
+  }
+
+  // (X && Y) || (X && Z) -> X && (Y || Z).
+  if (match(&R, m_BinaryOr(m_LogicalAnd(m_VPValue(X), m_VPValue(Y)),
+                           m_LogicalAnd(m_VPValue(X1), m_VPValue(Z)))) &&
+      X == X1) {
+    VPBuilder Builder(&R);
+    VPInstruction *YorZ = Builder.createOr(Y, Z, R.getDebugLoc());
+    VPInstruction *VPI = Builder.createLogicalAnd(X, YorZ, R.getDebugLoc());
+    R.getVPSingleValue()->replaceAllUsesWith(VPI);
+    // Order of simplification matters: simplify sub-recipes before root
+    // recipes.
+    return {YorZ, VPI};
+  }
+
+  return {};
 }
 
 /// Try to simplify the recipes in \p Plan.
@@ -1013,10 +1064,14 @@ static void simplifyRecipes(VPlan &Plan, LLVMContext &Ctx) {
       Plan.getEntry());
   VPTypeAnalysis TypeInfo(Plan.getCanonicalIV()->getScalarType(), Ctx);
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
+    std::deque<VPRecipeBase *> Worklist;
     for (auto &R : make_early_inc_range(*VPBB)) {
-      VPRecipeBase *NewR = simplifyRecipe(R, TypeInfo);
-      while (NewR)
-        NewR = simplifyRecipe(*NewR, TypeInfo);
+      Worklist.emplace_front(&R);
+      while (!Worklist.empty()) {
+        VPRecipeBase *R = Worklist.front();
+        Worklist.pop_front();
+        append_range(Worklist, simplifyRecipe(*R, TypeInfo, Plan));
+      }
     }
   }
 }