- Removed Multiplier from existing ActiveLaneMask VPInstruction

- Moved ICMP optimisation to optimizeForVFAndUF
- Remove preferTailFoldedLoop from selectInterleaveCount()
- Add test with i32 trip count

Author: kmclaughlin-arm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -3633,7 +3633,7 @@ LoopVectorizationPlanner::selectInterleaveCount(VPlan &Plan, ElementCount VF,
   // Do not interleave tail-folded loops, as the overhead of multiple
   // instructions to calculate the predicate is likely not beneficial.
   // If an epilogue is not allowed for any other reason, do not interleave.
-  if (!CM.isEpilogueAllowed() || CM.preferTailFoldedLoop())
+  if (!CM.isEpilogueAllowed())
     return 1;
 
   if (any_of(Plan.getVectorLoopRegion()->getEntryBasicBlock()->phis(),

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1232,8 +1232,13 @@ class LLVM_ABI_FOR_TEST VPInstruction : public VPRecipeWithIRFlags,
     // Creates a mask where each lane is active (true) whilst the current
     // counter (first operand + index) is less than the second operand. i.e.
     //    mask[i] = icmpt ult (op0 + i), op1
-    // The size of the mask returned is VF * Multiplier (UF, third op).
+    // The size of the mask returned is VF.
+    // When unrolled, ActiveLaneMask is duplicated.
     ActiveLaneMask,
+    // As above, but takes an additional operand (Multiplier). The size of
+    // the mask returned is VF * Multiplier (UF, op2).
+    // WideActiveLaneMask is used for control flow and is unrolled by widening,
+    // with one extract vector created per unroll part.
     WideActiveLaneMask,
     // Extracts each unrolled part of a (VF * UF) widened vector/mask.
     ExtractVectorForPart,

llvm/lib/Transforms/Vectorize/VPlanLowering.cpp

@@ -150,17 +150,15 @@ addVPLaneMaskPhiAndUpdateExitBranch(VPlan &Plan) {
   // TODO: Check if dropping the flags is needed.
   TopRegion->clearCanonicalIVNUW(CanonicalIVIncrement);
   DebugLoc DL = CanonicalIVIncrement->getDebugLoc();
-  // We can't use StartV directly in the WideActiveLaneMask
-  // VPInstruction, since we have to take unrolling into account.
-  // Each part needs to start at Part * VF
   auto *VecPreheader = Plan.getVectorPreheader();
   VPBuilder Builder(VecPreheader);
-
-  // Create the WideActiveLaneMask instruction using the correct
-  // start values.
   VPValue *TC = Plan.getTripCount();
   VPValue *VF = &Plan.getVF();
 
+  // We can't use StartV directly in the WideActiveLaneMask
+  // VPInstruction, since we have to take unrolling into account.
+  // CanonicalIVIncrementForPart is needed so that ExtractVectorForPart
+  // is unrolled for each part, rather than each extract starting at 0.
   auto *EntryIncrement =
       Builder.createOverflowingOp(VPInstruction::CanonicalIVIncrementForPart,
                                   {StartV, VF}, {}, DL, "index.part.next");
@@ -224,12 +222,9 @@ void VPlanTransforms::materializeHeaderMask(
       VPIRFlags::WrapFlagsTy(/*HasNUW=*/true, /*HasNSW=*/false)));
   VPValue *Mask;
   if (UseActiveLaneMask) {
-    VPValue *ALMMultiplier =
-        Plan.getConstantInt(LoopRegion->getCanonicalIVType(), 1);
-    Mask = Builder.createNaryOp(
-        VPInstruction::ActiveLaneMask,
-        {WideCanonicalIV, Plan.getTripCount(), ALMMultiplier}, nullptr,
-        "active.lane.mask");
+    Mask = Builder.createNaryOp(VPInstruction::ActiveLaneMask,
+                                {WideCanonicalIV, Plan.getTripCount()}, nullptr,
+                                "active.lane.mask");
   } else {
     Mask = Builder.createICmp(CmpInst::ICMP_ULE, WideCanonicalIV,
                               Plan.getOrCreateBackedgeTakenCount());

llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h

@@ -454,12 +454,6 @@ m_ExtractPenultimateElement(const Op0_t &Op0) {
   return m_VPInstruction<VPInstruction::ExtractPenultimateElement>(Op0);
 }
 
-template <typename Op0_t, typename Op1_t, typename Op2_t>
-inline VPInstruction_match<VPInstruction::ActiveLaneMask, Op0_t, Op1_t, Op2_t>
-m_ActiveLaneMask(const Op0_t &Op0, const Op1_t &Op1, const Op2_t &Op2) {
-  return m_VPInstruction<VPInstruction::ActiveLaneMask>(Op0, Op1, Op2);
-}
-
 template <typename Op0_t, typename Op1_t, typename Op2_t>
 inline VPInstruction_match<VPInstruction::WideActiveLaneMask, Op0_t, Op1_t,
                            Op2_t>

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -649,6 +649,7 @@ unsigned VPInstruction::getNumOperandsForOpcode() const {
   case Instruction::FCmp:
   case Instruction::ExtractElement:
   case Instruction::Store:
+  case VPInstruction::ActiveLaneMask:
   case VPInstruction::BranchOnCount:
   case VPInstruction::BranchOnTwoConds:
   case VPInstruction::FirstOrderRecurrenceSplice:
@@ -663,7 +664,6 @@ unsigned VPInstruction::getNumOperandsForOpcode() const {
     return 2;
   case Instruction::InsertElement:
   case Instruction::Select:
-  case VPInstruction::ActiveLaneMask:
   case VPInstruction::WideActiveLaneMask:
   case VPInstruction::ReductionStartVector:
     return 3;
@@ -795,9 +795,13 @@ Value *VPInstruction::generate(VPTransformState &State) {
     // Get the original loop tripcount.
     Value *ScalarTC = State.get(getOperand(1), VPLane(0));
 
+    unsigned Multiplier =
+        getOpcode() == VPInstruction::ActiveLaneMask
+            ? 1
+            : cast<VPConstantInt>(getOperand(2))->getZExtValue();
+
     // If this part of the active lane mask is scalar, generate the CMP directly
     // to avoid unnecessary extracts.
-    unsigned Multiplier = cast<VPConstantInt>(getOperand(2))->getZExtValue();
     if (State.VF.isScalar() && Multiplier == 1)
       return Builder.CreateCmp(CmpInst::Predicate::ICMP_ULT, VIVElem0, ScalarTC,
                                Name);
@@ -1116,23 +1120,8 @@ Value *VPInstruction::generate(VPTransformState &State) {
     if (Src->getType() == DstTy)
       return Src;
 
-    // If the VF is scalar & this is an extract of an active lane mask,
-    // generate an ICMP directly.
-    VPValue *Start, *TC;
-    if (State.VF.isScalar() &&
-        match(getOperand(0),
-              m_WideActiveLaneMask(m_VPValue(Start), m_VPValue(TC),
-                                   m_VPValue()))) {
-      Value *StartV = State.get(Start);
-      if (Part > 0)
-        StartV = Builder.CreateAdd(StartV, State.get(getOperand(1)));
-      return Builder.CreateCmp(CmpInst::Predicate::ICMP_ULT, StartV,
-                               State.get(TC));
-    }
-
-    auto *Idx = ConstantInt::get(Builder.getInt64Ty(),
-                                 State.VF.getKnownMinValue() * Part);
-    return Builder.CreateExtractVector(DstTy, Src, Idx);
+    return Builder.CreateExtractVector(
+        DstTy, Src, Builder.getInt64(State.VF.getKnownMinValue() * Part), Name);
   }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
@@ -1444,7 +1433,10 @@ InstructionCost VPInstruction::computeCost(ElementCount VF,
   case VPInstruction::ActiveLaneMask:
   case VPInstruction::WideActiveLaneMask: {
     Type *ArgTy = getOperand(0)->getScalarType();
-    unsigned Multiplier = cast<VPConstantInt>(getOperand(2))->getZExtValue();
+    unsigned Multiplier =
+        getOpcode() == VPInstruction::ActiveLaneMask
+            ? 1
+            : cast<VPConstantInt>(getOperand(2))->getZExtValue();
     Type *RetTy = toVectorTy(Type::getInt1Ty(Ctx.LLVMCtx), VF * Multiplier);
     IntrinsicCostAttributes Attrs(Intrinsic::get_active_lane_mask, RetTy,
                                   {ArgTy, ArgTy});

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2035,6 +2035,43 @@ static bool isConditionTrueViaVFAndUF(VPValue *Cond, VPlan &Plan,
   return SE.isKnownPredicate(CmpInst::ICMP_EQ, VectorTripCount, C);
 }
 
+static bool replaceMaskWithCompare(VPlan &Plan, ElementCount BestVF) {
+  if (!BestVF.isScalar())
+    return false;
+
+  bool MadeChange = false;
+  VPBuilder Builder;
+  VPRegionBlock *VectorRegion = Plan.getVectorLoopRegion();
+  VPBasicBlock *PreheaderVPBB = Plan.getVectorPreheader();
+  VPBasicBlock *ExitingVPBB = VectorRegion->getExitingBasicBlock();
+
+  VPValue *Start, *TC;
+  uint64_t Idx;
+  for (VPBasicBlock *VPBB : {PreheaderVPBB, ExitingVPBB}) {
+    for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
+      if (!match(&R, m_ExtractVectorForPart(
+                         m_WideActiveLaneMask(m_VPValue(Start), m_VPValue(TC),
+                                              m_VPValue()),
+                         m_ConstantInt(Idx))))
+        continue;
+
+      auto *Extract = cast<VPInstruction>(&R);
+      Builder.setInsertPoint(Extract);
+
+      if (Idx > 0)
+        Start = Builder.createAdd(
+            Start, Plan.getConstantInt(Start->getScalarType(), Idx));
+
+      VPValue *ICmp = Builder.createICmp(CmpInst::ICMP_ULT, Start, TC);
+      Extract->replaceAllUsesWith(ICmp);
+      Extract->eraseFromParent();
+      MadeChange = true;
+    }
+  }
+
+  return MadeChange;
+}
+
 /// Try to simplify the branch condition of \p Plan. This may restrict the
 /// resulting plan to \p BestVF and \p BestUF.
 static bool simplifyBranchConditionForVFAndUF(VPlan &Plan, ElementCount BestVF,
@@ -2050,16 +2087,13 @@ static bool simplifyBranchConditionForVFAndUF(VPlan &Plan, ElementCount BestVF,
   if (match(Term, m_BranchOnCount(
                       m_CombineOr(m_CanIVInc, m_c_Add(m_CanIVInc, m_LiveIn())),
                       m_VPValue())) ||
-      match(Term, m_BranchOnCond(m_Not(m_ActiveLaneMask(
-                      m_VPValue(), m_VPValue(), m_VPValue())))) ||
       match(Term,
             m_BranchOnCond(m_Not(m_ExtractVectorForPart(
                 m_WideActiveLaneMask(m_VPValue(), m_VPValue(), m_VPValue()),
                 m_ZeroInt()))))) {
     // Try to simplify the branch condition if VectorTC <= VF * UF when the
-    // latch terminator is BranchOnCount, BranchOnCond(Not(ActiveLaneMask)) or
-    // BranchOnCond(Not(ExtractVectorForPart(WideActiveLaneMask),
-    // 0))
+    // latch terminator is BranchOnCount or
+    // BranchOnCond(Not(ExtractVectorForPart(WideActiveLaneMask), 0))
     const SCEV *VectorTripCount =
         vputils::getSCEVExprForVPValue(&Plan.getVectorTripCount(), PSE);
     if (isa<SCEVCouldNotCompute>(VectorTripCount))
@@ -2103,8 +2137,8 @@ void VPlanTransforms::optimizeForVFAndUF(VPlan &Plan, ElementCount BestVF,
   assert(Plan.hasVF(BestVF) && "BestVF is not available in Plan");
   assert(Plan.hasUF(BestUF) && "BestUF is not available in Plan");
 
-  bool MadeChange =
-      simplifyBranchConditionForVFAndUF(Plan, BestVF, BestUF, PSE);
+  bool MadeChange = replaceMaskWithCompare(Plan, BestVF);
+  MadeChange |= simplifyBranchConditionForVFAndUF(Plan, BestVF, BestUF, PSE);
   MadeChange |= optimizeVectorInductionWidthForTCAndVFUF(Plan, BestVF, BestUF);
 
   if (MadeChange) {
@@ -3102,13 +3136,12 @@ static bool handleUncountableExitsWithSideEffects(
   VPBuilder MaskBuilder(HeaderVPBB, InsertIt);
   VPValue *FirstActive = MaskBuilder.createFirstActiveLane(*Cond);
   Type *IVScalarTy = IV->getScalarType();
-  VPValue *ALMMultiplier = Plan.getConstantInt(IVScalarTy, 1);
   VPValue *Zero = Plan.getZero(IVScalarTy);
   FirstActive =
       MaskBuilder.createScalarZExtOrTrunc(FirstActive, IVScalarTy, DebugLoc());
   VPValue *Mask = MaskBuilder.createNaryOp(VPInstruction::ActiveLaneMask,
-                                           {Zero, FirstActive, ALMMultiplier},
-                                           DebugLoc(), "uncountable.exit.mask");
+                                           {Zero, FirstActive}, DebugLoc(),
+                                           "uncountable.exit.mask");
 
   // Convert all other memory operations to use the mask.
   for (VPBasicBlock *VPBB : vp_rpo_plain_cfg_loop_body(HeaderVPBB))

llvm/lib/Transforms/Vectorize/VPlanUnroll.cpp

@@ -481,7 +481,7 @@ void UnrollState::unrollBlock(VPBlockBase *VPB) {
               m_WideActiveLaneMask(m_VPValue(), m_VPValue(), m_VPValue()))) {
       auto *ALM = cast<VPInstruction>(&R);
       addUniformForAllParts(ALM);
-      ALM->setOperand(2, Plan.getConstantInt(64, UF));
+      ALM->setOperand(2, getConstantInt(UF));
       continue;
     }
 

llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp

@@ -200,11 +200,6 @@ bool VPlanVerifier::verifyLastActiveLaneRecipe(
         match(Mask, m_VPInstruction<VPInstruction::WideActiveLaneMask>()))
       continue;
 
-    if (match(Mask, m_ActiveLaneMask(m_VPValue(), m_VPValue(), m_VPValue())) ||
-        match(Mask,
-              m_WideActiveLaneMask(m_VPValue(), m_VPValue(), m_VPValue())))
-      continue;
-
     CmpPredicate Pred;
     VPValue *LHS, *RHS;
     if (match(Mask, m_ICmp(Pred, m_VPValue(LHS), m_VPValue(RHS))) &&