[LoopVectorize] Introduce trip count minimal value threshold to ignore tail-folding for scalable vectors

[dtemirbulatov]

Summary

After D121595 was commited, I noticed regressions assosicated with small trip count numbers vectorization by tail folding with scalable vectors. As a solution for those issues I propose to introduce the minimal trip count threshold value.

[paulwalker-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

192

Typo here, should be TailFolding?

[dtemirbulatov]

Fixed typo.

[dtemirbulatov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

192

Done. thanks Paul.

[paulwalker-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

192–195

We're late into the LLVM 15 release cycle so here we're trying to fix an LLVM 14 regression without resorting to reversing the decision to use tail folding for all low trip count loops.

I think it'll be better to now do this in a way with the most minimal impact to other targets and so rather than a global command line option can you instead add a TTI hook. That way we can restrict the change to just AArch64/SVE and if it turns out we can rely on a more accurate cost model in the future and there are no other uses at that time, we can remove the TTI hook.

5078–5082

This doesn't look like the correct place to decide this. Within LoopVectorizePass::processLoop the code exists to choose CM_ScalarEpilogueNotAllowedLowTripLoop for low trip count looks. I think that's the code we want to change to also consider the new tail folding specific watermark.

[dtemirbulatov]

Fixed remarks.

[paulwalker-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

10119–10122

Can this be

if (*ExpectedTC > TTI->getMinTripCountTailFoldingThreshold())
  SEL = CM_ScalarEpilogueNotAllowedLowTripLoop

and still do what you need?

[dtemirbulatov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

10119–10122

I have to return false in order to prevent vectorization.

[paulwalker-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

10119–10122

Thanks @dtemirbulatov. In which case it looks like you need to restructure the code a little so we emit the necessary remark and debug information. e.g.

if (Hints.getForce() == LoopVectorizeHints::FK_Enabled)
...
else if (*ExpectedTC > TTI->getMinTripCountTailFoldingThreshold()) {
  LLVM_DEBUG(dbgs() << "\n");
  SEL = CM_ScalarEpilogueNotAllowedLowTripLoop;
} else {
  LLVM_DEBUG(dbgs() << " But the target considers the trip count too small to consider vectorizing.\n");
  reportVectorizationFailure(....);
  Hints.emitRemarkWithHints();
  return false;
}

[dtemirbulatov]

Updating proposed change according to @paulwalker-arm remarks.

[paulwalker-arm]

A minor pre-push change but otherwise looks good. Once landed can you please do the necessary to request an LLVM 15 backport and then we can turn our attention to potentially fixing up the cost model so this rather brute force fix can be removed.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

10123

To make the output consistent please add a <space> before the 'But' because otherwise you'll end up printing overheads are incurred.But the target considers...

[dtemirbulatov]

this was committed with cab6cd68340255be241b7cf169c67a1899ced115
but I made extra spaces before "Differential Revision", that is why it didn't close automatically.

llvm/include/llvm/Analysis/TargetTransformInfo.h

   /// \returns True if the target wants to expand the given reduction intrinsic
   /// into a shuffle sequence.
   bool shouldExpandReduction(const IntrinsicInst *II) const;
 
   /// \returns the size cost of rematerializing a GlobalValue address relative
   /// to a stack reload.
   unsigned getGISelRematGlobalCost() const;
 
+  /// \returns the lower bound of a trip count to decide on vectorization
+  /// while tail-folding.
+  unsigned getMinTripCountTailFoldingThreshold() const;
+
   /// \returns True if the target supports scalable vectors.
   bool supportsScalableVectors() const;
 
   /// \return true when scalable vectorization is preferred.
   bool enableScalableVectorization() const;
 
   /// \name Vector Predication Information
   /// @{
virtual unsigned getStoreVectorFactor(unsigned VF, unsigned StoreSize,
                                         unsigned ChainSizeInBytes,
                                         VectorType *VecTy) const = 0;
   virtual bool preferInLoopReduction(unsigned Opcode, Type *Ty,
                                      ReductionFlags) const = 0;
   virtual bool preferPredicatedReductionSelect(unsigned Opcode, Type *Ty,
                                                ReductionFlags) const = 0;
   virtual bool shouldExpandReduction(const IntrinsicInst *II) const = 0;
   virtual unsigned getGISelRematGlobalCost() const = 0;
+  virtual unsigned getMinTripCountTailFoldingThreshold() const = 0;
   virtual bool enableScalableVectorization() const = 0;
   virtual bool supportsScalableVectors() const = 0;
   virtual bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                                      Align Alignment) const = 0;
   virtual InstructionCost getInstructionLatency(const Instruction *I) = 0;
   virtual VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const = 0;
 };
public:
   bool shouldExpandReduction(const IntrinsicInst *II) const override {
     return Impl.shouldExpandReduction(II);
   }
 
   unsigned getGISelRematGlobalCost() const override {
     return Impl.getGISelRematGlobalCost();
   }
 
+  unsigned getMinTripCountTailFoldingThreshold() const override {
+    return Impl.getMinTripCountTailFoldingThreshold();
+  }
+
   bool supportsScalableVectors() const override {
     return Impl.supportsScalableVectors();
   }
 
   bool enableScalableVectorization() const override {
     return Impl.enableScalableVectorization();
   }
 

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

bool preferPredicatedReductionSelect(unsigned Opcode, Type *Ty,
                                        TTI::ReductionFlags Flags) const {
     return false;
   }
 
   bool shouldExpandReduction(const IntrinsicInst *II) const { return true; }
 
   unsigned getGISelRematGlobalCost() const { return 1; }
 
+  unsigned getMinTripCountTailFoldingThreshold() const { return 0; }
+
   bool supportsScalableVectors() const { return false; }
 
   bool enableScalableVectorization() const { return false; }
 
   bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                              Align Alignment) const {
     return false;
   }

llvm/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::shouldExpandReduction(const IntrinsicInst *II) const {
   return TTIImpl->shouldExpandReduction(II);
 }
 
 unsigned TargetTransformInfo::getGISelRematGlobalCost() const {
   return TTIImpl->getGISelRematGlobalCost();
 }
 
+unsigned TargetTransformInfo::getMinTripCountTailFoldingThreshold() const {
+  return TTIImpl->getMinTripCountTailFoldingThreshold();
+}
+
 bool TargetTransformInfo::supportsScalableVectors() const {
   return TTIImpl->supportsScalableVectors();
 }
 
 bool TargetTransformInfo::enableScalableVectorization() const {
   return TTIImpl->enableScalableVectorization();
 }
 

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

shouldConsiderAddressTypePromotion(const Instruction &I,
                                      bool &AllowPromotionWithoutCommonHeader);
 
   bool shouldExpandReduction(const IntrinsicInst *II) const { return false; }
 
   unsigned getGISelRematGlobalCost() const {
     return 2;
   }
 
+  unsigned getMinTripCountTailFoldingThreshold() const {
+    return ST->hasSVE() ? 5 : 0;
+  }
+
   PredicationStyle emitGetActiveLaneMask() const {
     if (ST->hasSVE())
       return PredicationStyle::DataAndControlFlow;
     return PredicationStyle::None;
   }
 
   bool preferPredicateOverEpilogue(Loop *L, LoopInfo *LI, ScalarEvolution &SE,
                                    AssumptionCache &AC, TargetLibraryInfo *TLI,

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

cl::desc("When epilogue vectorization is enabled, and a value greater than "
              "1 is specified, forces the given VF for all applicable epilogue "
              "loops."));
 
 static cl::opt<unsigned> EpilogueVectorizationMinVF(
     "epilogue-vectorization-minimum-VF", cl::init(16), cl::Hidden,
     cl::desc("Only loops with vectorization factor equal to or larger than "
              "the specified value are considered for epilogue vectorization."));
 
 /// Loops with a known constant trip count below this number are vectorized only
 /// if no scalar iteration overheads are incurred.
 static cl::opt<unsigned> TinyTripCountVectorThreshold(
     "vectorizer-min-trip-count", cl::init(16), cl::Hidden,
     cl::desc("Loops with a constant trip count that is smaller than this "
              "value are vectorized only if no scalar iteration overheads "
              "are incurred."));
 
 static cl::opt<unsigned> VectorizeMemoryCheckThreshold(
     "vectorize-memory-check-threshold", cl::init(128), cl::Hidden,
     cl::desc("The maximum allowed number of runtime memory checks"));
 
const SCEV *Rem = SE->getURemExpr(
         SE->getConstant(BackedgeTakenCount->getType(), MaxVFtimesIC));
     if (Rem->isZero()) {
       // Accept MaxFixedVF if we do not have a tail.
       LLVM_DEBUG(dbgs() << "LV: No tail will remain for any chosen VF.\n");
       return MaxFactors;
     }
   }
 
   // If we don't know the precise trip count, or if the trip count that we
   // found modulo the vectorization factor is not zero, try to fold the tail
   // by masking.
   // FIXME: look for a smaller MaxVF that does divide TC rather than masking.
   if (Legal->prepareToFoldTailByMasking()) {
     FoldTailByMasking = true;
     return MaxFactors;
   }
 
   // If there was a tail-folding hint/switch, but we can't fold the tail by
   // masking, fallback to a vectorization with a scalar epilogue.
   if (ScalarEpilogueStatus == CM_ScalarEpilogueNotNeededUsePredicate) {
     LLVM_DEBUG(dbgs() << "LV: Cannot fold tail by masking: vectorize with a "
#endif /* NDEBUG */
   if (ExpectedTC && *ExpectedTC < TinyTripCountVectorThreshold) {
     LLVM_DEBUG(dbgs() << "LV: Found a loop with a very small trip count. "
                       << "This loop is worth vectorizing only if no scalar "
                       << "iteration overheads are incurred.");
     if (Hints.getForce() == LoopVectorizeHints::FK_Enabled)
       LLVM_DEBUG(dbgs() << " But vectorizing was explicitly forced.\n");
     else {
       LLVM_DEBUG(dbgs() << "\n");
       SEL = CM_ScalarEpilogueNotAllowedLowTripLoop;
+      if (TTI->getMinTripCountTailFoldingThreshold() &&
+          ExpectedTC <= TTI->getMinTripCountTailFoldingThreshold())
+        return false;
     }
   }
 
   // Check the function attributes to see if implicit floats are allowed.
   // FIXME: This check doesn't seem possibly correct -- what if the loop is
   // an integer loop and the vector instructions selected are purely integer
   // vector instructions?
   if (F->hasFnAttribute(Attribute::NoImplicitFloat)) {
     reportVectorizationFailure(

llvm/test/Transforms/LoopVectorize/AArch64/sve-low-trip-count.ll

for.body:                                         ; preds = %entry, %for.body
   br i1 %exitcond.not, label %for.end, label %for.body
 
 for.end:                                          ; preds = %for.body
   ret void
 }
 
 define void @trip5_i8(i8* noalias nocapture noundef %dst, i8* noalias nocapture noundef readonly %src) #0 {
 ; CHECK-LABEL: @trip5_i8(
-; CHECK:       vector.body:
-; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %vector.body ]
-; CHECK:         [[ACTIVE_LANE_MASK:%.*]] = phi <vscale x 16 x i1> [ {{%.*}}, %vector.ph ], [ [[ACTIVE_LANE_MASK_NEXT:%.*]], %vector.body ]
-; CHECK:         {{%.*}} = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]], <vscale x 16 x i8> poison)
-; CHECK:         {{%.*}} = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]], <vscale x 16 x i8> poison)
-; CHECK:         call void @llvm.masked.store.nxv16i8.p0nxv16i8(<vscale x 16 x i8> {{%.*}}, <vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]])
-; CHECK:         [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    [[VF:%.*]] = mul i64 [[VSCALE]], 16
-; CHECK-NEXT:    [[INDEX_NEXT]] = add i64 [[INDEX]], [[VF]]
-; CHECK-NEXT:    [[ACTIVE_LANE_MASK_NEXT]] = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 [[INDEX_NEXT]], i64 5)
-; CHECK-NEXT:    [[ACTIVE_LANE_MASK_NOT:%.*]] = xor <vscale x 16 x i1> [[ACTIVE_LANE_MASK_NEXT]], shufflevector (<vscale x 16 x i1> insertelement (<vscale x 16 x i1> poison, i1 true, i32 0), <vscale x 16 x i1> poison, <vscale x 16 x i32> zeroinitializer)
-; CHECK-NEXT:    br i1 true, label %middle.block, label %vector.body
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_08:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i8, i8* [[SRC:%.*]], i64 [[I_08]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, i8* [[ARRAYIDX]], align 1
+; CHECK-NEXT:    [[MUL:%.*]] = shl i8 [[TMP0]], 1
+; CHECK-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i8, i8* [[DST:%.*]], i64 [[I_08]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load i8, i8* [[ARRAYIDX1]], align 1
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[MUL]], [[TMP1]]
+; CHECK-NEXT:    store i8 [[ADD]], i8* [[ARRAYIDX1]], align 1
+; CHECK-NEXT:    [[INC]] = add nuw nsw i64 [[I_08]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[INC]], 5
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
+; CHECK:       for.end:
+; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body
   %i.08 = phi i64 [ 0, %entry ], [ %inc, %for.body ]
   %arrayidx = getelementptr inbounds i8, i8* %src, i64 %i.08
   %0 = load i8, i8* %arrayidx, align 1