[llvm] [LV] Enable partial reductions with EVL tail folding (PR #215515)
Elvis Wang via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 11 19:15:45 PDT 2026
================
@@ -0,0 +1,117 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals none --filter-out-after "^scalar.ph:" --version 4
+; RUN: opt -passes=loop-vectorize -mattr=+v,+experimental-zvdot4a8i -tail-folding-policy=dont-fold-tail -S < %s | FileCheck %s --check-prefixes=NOTAILFOLD
+; RUN: opt -passes=loop-vectorize -mattr=+v,+experimental-zvdot4a8i -S < %s | FileCheck %s --check-prefixes=TAILFOLD
+
+; A data-dependent predicated dot-product reduction. This used to crash the
+; loop vectorizer under EVL tail folding (llvm/llvm-project#167861): the scaled
+; reduction PHI was fed by a plain vp.reduce.add of a mismatched type instead
+; of a partial reduction. Check that a partial reduction is formed with the
+; predication folded into the reduction input, both with and without EVL tail
+; folding.
+
+target triple = "riscv64-none-unknown-elf"
+
+define i32 @pred_dot(ptr %a, ptr %b, ptr %c, i64 %n) {
+; NOTAILFOLD-LABEL: define i32 @pred_dot(
+; NOTAILFOLD-SAME: ptr [[A:%.*]], ptr [[B:%.*]], ptr [[C:%.*]], i64 [[N:%.*]]) #[[ATTR0:[0-9]+]] {
+; NOTAILFOLD-NEXT: entry:
+; NOTAILFOLD-NEXT: [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; NOTAILFOLD-NEXT: [[TMP1:%.*]] = shl nuw i64 [[TMP0]], 2
+; NOTAILFOLD-NEXT: [[TMP2:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP1]], i64 8)
+; NOTAILFOLD-NEXT: [[MIN_ITERS_CHECK:%.*]] = icmp ult i64 [[N]], [[TMP2]]
+; NOTAILFOLD-NEXT: br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; NOTAILFOLD: vector.ph:
+; NOTAILFOLD-NEXT: [[TMP3:%.*]] = shl nuw i64 [[TMP0]], 2
+; NOTAILFOLD-NEXT: [[N_MOD_VF:%.*]] = urem i64 [[N]], [[TMP3]]
+; NOTAILFOLD-NEXT: [[N_VEC:%.*]] = sub i64 [[N]], [[N_MOD_VF]]
+; NOTAILFOLD-NEXT: br label [[VECTOR_BODY:%.*]]
+; NOTAILFOLD: vector.body:
+; NOTAILFOLD-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; NOTAILFOLD-NEXT: [[VEC_PHI:%.*]] = phi <vscale x 1 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; NOTAILFOLD-NEXT: [[TMP4:%.*]] = getelementptr inbounds i8, ptr [[C]], i64 [[INDEX]]
+; NOTAILFOLD-NEXT: [[WIDE_LOAD:%.*]] = load <vscale x 4 x i8>, ptr [[TMP4]], align 1
+; NOTAILFOLD-NEXT: [[TMP5:%.*]] = icmp ne <vscale x 4 x i8> [[WIDE_LOAD]], zeroinitializer
+; NOTAILFOLD-NEXT: [[TMP6:%.*]] = getelementptr i8, ptr [[A]], i64 [[INDEX]]
+; NOTAILFOLD-NEXT: [[WIDE_MASKED_LOAD:%.*]] = call <vscale x 4 x i8> @llvm.masked.load.nxv4i8.p0(ptr align 1 [[TMP6]], <vscale x 4 x i1> [[TMP5]], <vscale x 4 x i8> poison)
+; NOTAILFOLD-NEXT: [[TMP7:%.*]] = getelementptr i8, ptr [[B]], i64 [[INDEX]]
+; NOTAILFOLD-NEXT: [[WIDE_MASKED_LOAD1:%.*]] = call <vscale x 4 x i8> @llvm.masked.load.nxv4i8.p0(ptr align 1 [[TMP7]], <vscale x 4 x i1> [[TMP5]], <vscale x 4 x i8> poison)
+; NOTAILFOLD-NEXT: [[TMP8:%.*]] = sext <vscale x 4 x i8> [[WIDE_MASKED_LOAD]] to <vscale x 4 x i32>
+; NOTAILFOLD-NEXT: [[TMP9:%.*]] = sext <vscale x 4 x i8> [[WIDE_MASKED_LOAD1]] to <vscale x 4 x i32>
+; NOTAILFOLD-NEXT: [[TMP10:%.*]] = mul nsw <vscale x 4 x i32> [[TMP8]], [[TMP9]]
+; NOTAILFOLD-NEXT: [[TMP11:%.*]] = select <vscale x 4 x i1> [[TMP5]], <vscale x 4 x i32> [[TMP10]], <vscale x 4 x i32> zeroinitializer
+; NOTAILFOLD-NEXT: [[PARTIAL_REDUCE]] = call <vscale x 1 x i32> @llvm.vector.partial.reduce.add.nxv1i32.nxv4i32(<vscale x 1 x i32> [[VEC_PHI]], <vscale x 4 x i32> [[TMP11]])
+; NOTAILFOLD-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], [[TMP3]]
+; NOTAILFOLD-NEXT: [[TMP12:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; NOTAILFOLD-NEXT: br i1 [[TMP12]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; NOTAILFOLD: middle.block:
+; NOTAILFOLD-NEXT: [[TMP13:%.*]] = call i32 @llvm.vector.reduce.add.nxv1i32(<vscale x 1 x i32> [[PARTIAL_REDUCE]])
+; NOTAILFOLD-NEXT: [[CMP_N:%.*]] = icmp eq i64 [[N]], [[N_VEC]]
+; NOTAILFOLD-NEXT: br i1 [[CMP_N]], label [[FOR_EXIT:%.*]], label [[SCALAR_PH]]
+; NOTAILFOLD: scalar.ph:
+;
+; TAILFOLD-LABEL: define i32 @pred_dot(
+; TAILFOLD-SAME: ptr [[A:%.*]], ptr [[B:%.*]], ptr [[C:%.*]], i64 [[N:%.*]]) #[[ATTR0:[0-9]+]] {
+; TAILFOLD-NEXT: entry:
+; TAILFOLD-NEXT: br label [[VECTOR_PH:%.*]]
+; TAILFOLD: vector.ph:
+; TAILFOLD-NEXT: br label [[VECTOR_BODY:%.*]]
+; TAILFOLD: vector.body:
+; TAILFOLD-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[CURRENT_ITERATION_NEXT:%.*]], [[VECTOR_BODY]] ]
+; TAILFOLD-NEXT: [[VEC_PHI:%.*]] = phi <vscale x 1 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; TAILFOLD-NEXT: [[AVL:%.*]] = phi i64 [ [[N]], [[VECTOR_PH]] ], [ [[AVL_NEXT:%.*]], [[VECTOR_BODY]] ]
+; TAILFOLD-NEXT: [[TMP0:%.*]] = call i32 @llvm.experimental.get.vector.length.i64(i64 [[AVL]], i32 4, i1 true)
+; TAILFOLD-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[C]], i64 [[INDEX]]
+; TAILFOLD-NEXT: [[VP_OP_LOAD:%.*]] = call <vscale x 4 x i8> @llvm.vp.load.nxv4i8.p0(ptr align 1 [[TMP1]], <vscale x 4 x i1> splat (i1 true), i32 [[TMP0]])
+; TAILFOLD-NEXT: [[TMP2:%.*]] = icmp ne <vscale x 4 x i8> [[VP_OP_LOAD]], zeroinitializer
+; TAILFOLD-NEXT: [[TMP3:%.*]] = call <vscale x 4 x i1> @llvm.vp.merge.nxv4i1(<vscale x 4 x i1> splat (i1 true), <vscale x 4 x i1> [[TMP2]], <vscale x 4 x i1> zeroinitializer, i32 [[TMP0]])
+; TAILFOLD-NEXT: [[TMP4:%.*]] = getelementptr i8, ptr [[A]], i64 [[INDEX]]
+; TAILFOLD-NEXT: [[VP_OP_LOAD1:%.*]] = call <vscale x 4 x i8> @llvm.vp.load.nxv4i8.p0(ptr align 1 [[TMP4]], <vscale x 4 x i1> [[TMP2]], i32 [[TMP0]])
+; TAILFOLD-NEXT: [[TMP5:%.*]] = getelementptr i8, ptr [[B]], i64 [[INDEX]]
+; TAILFOLD-NEXT: [[VP_OP_LOAD2:%.*]] = call <vscale x 4 x i8> @llvm.vp.load.nxv4i8.p0(ptr align 1 [[TMP5]], <vscale x 4 x i1> [[TMP2]], i32 [[TMP0]])
+; TAILFOLD-NEXT: [[TMP6:%.*]] = sext <vscale x 4 x i8> [[VP_OP_LOAD1]] to <vscale x 4 x i32>
+; TAILFOLD-NEXT: [[TMP7:%.*]] = sext <vscale x 4 x i8> [[VP_OP_LOAD2]] to <vscale x 4 x i32>
+; TAILFOLD-NEXT: [[TMP8:%.*]] = mul nsw <vscale x 4 x i32> [[TMP6]], [[TMP7]]
+; TAILFOLD-NEXT: [[TMP9:%.*]] = select <vscale x 4 x i1> [[TMP3]], <vscale x 4 x i32> [[TMP8]], <vscale x 4 x i32> zeroinitializer
+; TAILFOLD-NEXT: [[PARTIAL_REDUCE]] = call <vscale x 1 x i32> @llvm.vector.partial.reduce.add.nxv1i32.nxv4i32(<vscale x 1 x i32> [[VEC_PHI]], <vscale x 4 x i32> [[TMP9]])
----------------
ElvisWang123 wrote:
We should generate vp.merge instead of select.
https://github.com/llvm/llvm-project/pull/215515
More information about the llvm-commits
mailing list