[llvm] [LV] Fix a crash on SELECT masks with a one element vector. (PR #83332)

Cameron McInally via llvm-commits llvm-commits at lists.llvm.org
Wed Feb 28 13:06:16 PST 2024


https://github.com/mcinally created https://github.com/llvm/llvm-project/pull/83332

When VPlan generates a SELECT with a <1 x i1> mask, make sure to access part 0 of the condition.

Fixes llvm#73894.

>From e755d50e8663ece6fb71a278d5bce4cde04b6322 Mon Sep 17 00:00:00 2001
From: Cameron McInally <cmcinally at nvidia.com>
Date: Wed, 28 Feb 2024 12:46:55 -0800
Subject: [PATCH] [LV] Fix a crash on SELECT masks with a one element vector.

When VPlan generates a SELECT with a <1 x i1> mask, make sure to
access part 0 of the condition.

Fixes llvm#73894.
---
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp |   4 +-
 llvm/test/Transforms/LoopVectorize/pr73894.ll | 116 ++++++++++++++++++
 2 files changed, 119 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/pr73894.ll

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 2d2f6acf913f10..ba6bf22f85ffd9 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -302,7 +302,9 @@ Value *VPInstruction::generateInstruction(VPTransformState &State,
     return Builder.CreateCmp(getPredicate(), A, B, Name);
   }
   case Instruction::Select: {
-    Value *Cond = State.get(getOperand(0), Part);
+    Value *Cond = State.VF.isVector()
+                      ? State.get(getOperand(0), Part)
+                      : State.get(getOperand(0), VPIteration(Part, 0));
     Value *Op1 = State.get(getOperand(1), Part);
     Value *Op2 = State.get(getOperand(2), Part);
     return Builder.CreateSelect(Cond, Op1, Op2, Name);
diff --git a/llvm/test/Transforms/LoopVectorize/pr73894.ll b/llvm/test/Transforms/LoopVectorize/pr73894.ll
new file mode 100644
index 00000000000000..285a2e03415b51
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/pr73894.ll
@@ -0,0 +1,116 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=loop-vectorize -mcpu=neoverse-v1 -force-vector-interleave=2 -force-vector-width=1 -S %s | FileCheck %s
+
+target triple = "arm64-linux"
+
+define i32 @pr70988() {
+; CHECK-LABEL: define i32 @pr70988(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr null, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 15
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.umax.i32(i32 [[TMP1]], i32 1)
+; CHECK-NEXT:    [[UMAX:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    br i1 false, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[N_RND_UP:%.*]] = add i64 [[UMAX]], 1
+; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 [[N_RND_UP]], 2
+; CHECK-NEXT:    [[N_VEC:%.*]] = sub i64 [[N_RND_UP]], [[N_MOD_VF]]
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK_ENTRY:%.*]] = call <1 x i1> @llvm.get.active.lane.mask.v1i1.i64(i64 0, i64 [[UMAX]])
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK_ENTRY1:%.*]] = call <1 x i1> @llvm.get.active.lane.mask.v1i1.i64(i64 1, i64 [[UMAX]])
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT6:%.*]], [[PRED_LOAD_CONTINUE5:%.*]] ]
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK:%.*]] = phi <1 x i1> [ [[ACTIVE_LANE_MASK_ENTRY]], [[VECTOR_PH]] ], [ [[ACTIVE_LANE_MASK_NEXT:%.*]], [[PRED_LOAD_CONTINUE5]] ]
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK2:%.*]] = phi <1 x i1> [ [[ACTIVE_LANE_MASK_ENTRY1]], [[VECTOR_PH]] ], [ [[ACTIVE_LANE_MASK_NEXT7:%.*]], [[PRED_LOAD_CONTINUE5]] ]
+; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi i32 [ 0, [[VECTOR_PH]] ], [ [[TMP20:%.*]], [[PRED_LOAD_CONTINUE5]] ]
+; CHECK-NEXT:    [[VEC_PHI3:%.*]] = phi i32 [ 0, [[VECTOR_PH]] ], [ [[TMP22:%.*]], [[PRED_LOAD_CONTINUE5]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <1 x i1> [[ACTIVE_LANE_MASK]], i32 0
+; CHECK-NEXT:    br i1 [[TMP3]], label [[PRED_LOAD_IF:%.*]], label [[PRED_LOAD_CONTINUE:%.*]]
+; CHECK:       pred.load.if:
+; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[INDEX]], 0
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr i32, ptr null, i64 [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = load ptr, ptr [[TMP5]], align 8
+; CHECK-NEXT:    [[TMP7:%.*]] = load i32, ptr [[TMP6]], align 4
+; CHECK-NEXT:    br label [[PRED_LOAD_CONTINUE]]
+; CHECK:       pred.load.continue:
+; CHECK-NEXT:    [[TMP8:%.*]] = phi ptr [ poison, [[VECTOR_BODY]] ], [ [[TMP6]], [[PRED_LOAD_IF]] ]
+; CHECK-NEXT:    [[TMP9:%.*]] = phi i32 [ poison, [[VECTOR_BODY]] ], [ [[TMP7]], [[PRED_LOAD_IF]] ]
+; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <1 x i1> [[ACTIVE_LANE_MASK2]], i32 0
+; CHECK-NEXT:    br i1 [[TMP10]], label [[PRED_LOAD_IF4:%.*]], label [[PRED_LOAD_CONTINUE5]]
+; CHECK:       pred.load.if4:
+; CHECK-NEXT:    [[TMP11:%.*]] = add i64 [[INDEX]], 1
+; CHECK-NEXT:    [[TMP12:%.*]] = getelementptr i32, ptr null, i64 [[TMP11]]
+; CHECK-NEXT:    [[TMP13:%.*]] = load ptr, ptr [[TMP12]], align 8
+; CHECK-NEXT:    [[TMP14:%.*]] = load i32, ptr [[TMP13]], align 4
+; CHECK-NEXT:    br label [[PRED_LOAD_CONTINUE5]]
+; CHECK:       pred.load.continue5:
+; CHECK-NEXT:    [[TMP15:%.*]] = phi ptr [ poison, [[PRED_LOAD_CONTINUE]] ], [ [[TMP13]], [[PRED_LOAD_IF4]] ]
+; CHECK-NEXT:    [[TMP16:%.*]] = phi i32 [ poison, [[PRED_LOAD_CONTINUE]] ], [ [[TMP14]], [[PRED_LOAD_IF4]] ]
+; CHECK-NEXT:    [[TMP17:%.*]] = tail call i32 @llvm.smax.i32(i32 [[TMP9]], i32 [[VEC_PHI]])
+; CHECK-NEXT:    [[TMP18:%.*]] = tail call i32 @llvm.smax.i32(i32 [[TMP16]], i32 [[VEC_PHI3]])
+; CHECK-NEXT:    [[TMP19:%.*]] = extractelement <1 x i1> [[ACTIVE_LANE_MASK]], i32 0
+; CHECK-NEXT:    [[TMP20]] = select i1 [[TMP19]], i32 [[TMP17]], i32 [[VEC_PHI]]
+; CHECK-NEXT:    [[TMP21:%.*]] = extractelement <1 x i1> [[ACTIVE_LANE_MASK2]], i32 0
+; CHECK-NEXT:    [[TMP22]] = select i1 [[TMP21]], i32 [[TMP18]], i32 [[VEC_PHI3]]
+; CHECK-NEXT:    [[INDEX_NEXT:%.*]] = add i64 [[INDEX]], 2
+; CHECK-NEXT:    [[INDEX_NEXT6]] = add i64 [[INDEX]], 2
+; CHECK-NEXT:    [[TMP23:%.*]] = add i64 [[INDEX_NEXT]], 1
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK_NEXT]] = call <1 x i1> @llvm.get.active.lane.mask.v1i1.i64(i64 [[INDEX_NEXT]], i64 [[UMAX]])
+; CHECK-NEXT:    [[ACTIVE_LANE_MASK_NEXT7]] = call <1 x i1> @llvm.get.active.lane.mask.v1i1.i64(i64 [[TMP23]], i64 [[UMAX]])
+; CHECK-NEXT:    [[TMP24:%.*]] = xor <1 x i1> [[ACTIVE_LANE_MASK_NEXT]], <i1 true>
+; CHECK-NEXT:    [[TMP25:%.*]] = xor <1 x i1> [[ACTIVE_LANE_MASK_NEXT7]], <i1 true>
+; CHECK-NEXT:    [[TMP26:%.*]] = extractelement <1 x i1> [[TMP24]], i32 0
+; CHECK-NEXT:    br i1 [[TMP26]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[RDX_MINMAX:%.*]] = call i32 @llvm.smax.i32(i32 [[TMP20]], i32 [[TMP22]])
+; CHECK-NEXT:    br i1 true, label [[EXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[BC_MERGE_RDX:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[RDX_MINMAX]], [[MIDDLE_BLOCK]] ]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDUC:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[INDUC_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[MAX:%.*]] = phi i32 [ [[BC_MERGE_RDX]], [[SCALAR_PH]] ], [ [[TMP29:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, ptr null, i64 [[INDUC]]
+; CHECK-NEXT:    [[TMP27:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[TMP28:%.*]] = load i32, ptr [[TMP27]], align 4
+; CHECK-NEXT:    [[TMP29]] = tail call i32 @llvm.smax.i32(i32 [[TMP28]], i32 [[MAX]])
+; CHECK-NEXT:    [[INDUC_NEXT]] = add nuw nsw i64 [[INDUC]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[INDUC_NEXT]], [[UMAX]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[EXIT]], label [[LOOP]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RES:%.*]] = phi i32 [ [[TMP29]], [[LOOP]] ], [ [[RDX_MINMAX]], [[MIDDLE_BLOCK]] ]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+entry:
+  %0 = load i32, ptr null
+  %1 = and i32 %0, 15
+  %2 = call i32 @llvm.umax.i32(i32 %1, i32 1)
+  %umax = zext i32 %2 to i64
+  br label %loop
+
+loop:
+  %induc = phi i64 [ 0, %entry ], [ %induc.next, %loop ]
+  %max = phi i32 [ 0, %entry ], [ %5, %loop ]
+  %gep = getelementptr i32, ptr null, i64 %induc
+  %3 = load ptr, ptr %gep
+  %4 = load i32, ptr %3
+  %5 = tail call i32 @llvm.smax.i32(i32 %4, i32 %max)
+  %induc.next = add nuw nsw i64 %induc, 1
+  %exitcond.not = icmp eq i64 %induc.next, %umax
+  br i1 %exitcond.not, label %exit, label %loop
+
+exit:
+  %res = phi i32 [ %5, %loop ]
+  ret i32 %res
+}
+
+declare i32 @llvm.smax.i32(i32, i32)
+declare i32 @llvm.umax.i32(i32, i32)
+;.
+; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]], [[META2:![0-9]+]]}
+; CHECK: [[META1]] = !{!"llvm.loop.isvectorized", i32 1}
+; CHECK: [[META2]] = !{!"llvm.loop.unroll.runtime.disable"}
+; CHECK: [[LOOP3]] = distinct !{[[LOOP3]], [[META1]]}
+;.



More information about the llvm-commits mailing list