[LoopVectorize] Permit tail-folding for low trip counts using scalable vectors

[david-arm]

Summary

When the loop vectoriser encounters a known low trip count it tries
to create a single predicated loop in order to get the benefit of
vectorisation and eliminate the scalar tail. However, until now the
vectoriser prevented the use of scalable vectors in this case due
to concerns in the past about stability. I believe that tail-folded
loops using scalable vectors are now sufficiently well tested that
we can enable this. For the same reason I've also enabled it when
optimising for code size too.

Tests added here:

Transforms/LoopVectorize/AArch64/sve-low-trip-count.ll
Transforms/LoopVectorize/AArch64/sve-tail-folding-optsize.ll
Transforms/LoopVectorize/RISCV/low-trip-count.ll

[david-arm]

• Added more CHECK lines to the new tests because they will be useful for a future patch.

[sdesmalen]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5249–5250

I think we can remove the condition entirely, so that we consider using scalable vectors + tail folding when optimising for code-size as well.

I know that for SVE we'll want to improve code quality to avoid the redundant compare, but when optimising for code-size the user has made the decision that code-size is more important than performance. And the cost-model will still have a say in which is more beneficial (scalar, fixed or scalable) and may still choose a fixed-width VF in case the ScalableVF may not be legal for the loop.

[david-arm]

• Completely removed all restrictions on using tail-folding for scalable vectors.
• Added test to show we apply tail-folding when compiling with -Os

[david-arm]

Gentle ping.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

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;
     }
   }
 
-  // For scalable vectors don't use tail folding for low trip counts or
-  // optimizing for code size. We only permit this if the user has explicitly
-  // requested it.
-  if (ScalarEpilogueStatus != CM_ScalarEpilogueNotNeededUsePredicate &&
-      ScalarEpilogueStatus != CM_ScalarEpilogueNotAllowedUsePredicate &&
+  // For scalable vectors don't use tail folding when optimizing for code size.
+  if (ScalarEpilogueStatus == CM_ScalarEpilogueNotAllowedOptSize &&
       MaxFactors.ScalableVF.isVector())
     MaxFactors.ScalableVF = ElementCount::getScalable(0);
 
   // 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()) {

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

+; RUN: opt -loop-vectorize -S < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define void @trip7_i64(i64* noalias nocapture noundef %dst, i64* noalias nocapture noundef readonly %src) #0 {
+; CHECK-LABEL: @trip7_i64(
+; CHECK:       vector.body:
+; CHECK:         [[ACTIVE_LANE_MASK:%.*]] = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i64(i64 {{%.*}}, i64 7)
+; CHECK:         {{%.*}} = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64(<vscale x 2 x i64>* {{%.*}}, i32 8, <vscale x 2 x i1> [[ACTIVE_LANE_MASK]], <vscale x 2 x i64> poison)
+; CHECK:         {{%.*}} = call <vscale x 2 x i64> @llvm.masked.load.nxv2i64.p0nxv2i64(<vscale x 2 x i64>* {{%.*}}, i32 8, <vscale x 2 x i1> [[ACTIVE_LANE_MASK]], <vscale x 2 x i64> poison)
+; CHECK:         call void @llvm.masked.store.nxv2i64.p0nxv2i64(<vscale x 2 x i64> {{%.*}}, <vscale x 2 x i64>* {{%.*}}, i32 8, <vscale x 2 x i1> [[ACTIVE_LANE_MASK]])
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i.06 = phi i64 [ 0, %entry ], [ %inc, %for.body ]
+  %arrayidx = getelementptr inbounds i64, i64* %src, i64 %i.06
+  %0 = load i64, i64* %arrayidx, align 8
+  %mul = shl nsw i64 %0, 1
+  %arrayidx1 = getelementptr inbounds i64, i64* %dst, i64 %i.06
+  %1 = load i64, i64* %arrayidx1, align 8
+  %add = add nsw i64 %1, %mul
+  store i64 %add, i64* %arrayidx1, align 8
+  %inc = add nuw nsw i64 %i.06, 1
+  %exitcond.not = icmp eq i64 %inc, 7
+  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:         [[ACTIVE_LANE_MASK:%.*]] = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 {{%.*}}, i64 5)
+; 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]])
+;
+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
+  %mul = shl i8 %0, 1
+  %arrayidx1 = getelementptr inbounds i8, i8* %dst, i64 %i.08
+  %1 = load i8, i8* %arrayidx1, align 1
+  %add = add i8 %mul, %1
+  store i8 %add, i8* %arrayidx1, align 1
+  %inc = add nuw nsw i64 %i.08, 1
+  %exitcond.not = icmp eq i64 %inc, 5
+  br i1 %exitcond.not, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
+attributes #0 = { vscale_range(1,16) "target-features"="+sve" }

llvm/test/Transforms/LoopVectorize/RISCV/low-trip-count.ll

+; RUN: opt -loop-vectorize -riscv-v-vector-bits-min=128 -scalable-vectorization=on -force-target-instruction-cost=1 -S < %s | FileCheck %s
+
+target triple = "riscv64"
+
+define void @trip5_i8(i8* noalias nocapture noundef %dst, i8* noalias nocapture noundef readonly %src) #0 {
+; CHECK-LABEL: @trip5_i8(
+; CHECK:       vector.body:
+; CHECK:         [[ACTIVE_LANE_MASK:%.*]] = icmp ule <vscale x 8 x i64> {{%.*}}, shufflevector (<vscale x 8 x i64> insertelement (<vscale x 8 x i64> poison, i64 4, i32 0), <vscale x 8 x i64> poison, <vscale x 8 x i32> zeroinitializer)
+; CHECK:         {{%.*}} = call <vscale x 8 x i8> @llvm.masked.load.nxv8i8.p0nxv8i8(<vscale x 8 x i8>* {{%.*}}, i32 1, <vscale x 8 x i1> [[ACTIVE_LANE_MASK]], <vscale x 8 x i8> poison)
+; CHECK:         {{%.*}} = call <vscale x 8 x i8> @llvm.masked.load.nxv8i8.p0nxv8i8(<vscale x 8 x i8>* {{%.*}}, i32 1, <vscale x 8 x i1> [[ACTIVE_LANE_MASK]], <vscale x 8 x i8> poison)
+; CHECK:         call void @llvm.masked.store.nxv8i8.p0nxv8i8(<vscale x 8 x i8> {{%.*}}, <vscale x 8 x i8>* {{%.*}}, i32 1, <vscale x 8 x i1> [[ACTIVE_LANE_MASK]])
+;
+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
+  %mul = shl i8 %0, 1
+  %arrayidx1 = getelementptr inbounds i8, i8* %dst, i64 %i.08
+  %1 = load i8, i8* %arrayidx1, align 1
+  %add = add i8 %mul, %1
+  store i8 %add, i8* %arrayidx1, align 1
+  %inc = add nuw nsw i64 %i.08, 1
+  %exitcond.not = icmp eq i64 %inc, 5
+  br i1 %exitcond.not, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
+attributes #0 = { "target-features"="+v,+d" }