[VPlan] Fix assert in store-user in narrowToSingleScalars

[artagnon]

Follow up on c2d4c7c ([VPlan] Permit more users in narrowToSingleScalars) to fix an assert related to WidenStore users of the recipe being narrowed in narrowToSingleScalars.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Ramkumar Ramachandra (artagnon)

Changes

Follow up on c2d4c7c ([VPlan] Permit more users in narrowToSingleScalars) to fix an assert related to WidenStore users of the recipe being narrowed in narrowToSingleScalars.

---

Full diff: https://github.com/llvm/llvm-project/pull/167686.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+1)
• (added) llvm/test/Transforms/LoopVectorize/X86/narrow-to-single-scalar.ll (+53)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 7cef98f465715..bf9cffc5dfe31 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1437,6 +1437,7 @@ static void narrowToSingleScalarRecipes(VPlan &Plan) {
               // we checked the RepOrWidenR operand against
               // vputils::isSingleScalar.
               assert(RepOrWidenR == Store->getAddr() ||
+                     RepOrWidenR == Store->getMask() ||
                      vputils::isSingleScalar(Store->getStoredValue()));
               return true;
             }
diff --git a/llvm/test/Transforms/LoopVectorize/X86/narrow-to-single-scalar.ll b/llvm/test/Transforms/LoopVectorize/X86/narrow-to-single-scalar.ll
new file mode 100644
index 0000000000000..94a05a67a0bdc
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/X86/narrow-to-single-scalar.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -p loop-vectorize -mcpu=skylake -S %s | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@p = external global [3952 x i8], align 8
+@q = external global [3952 x i8], align 8
+
+define void @narrow_store_user_mask_operand(i32 %x) {
+; CHECK-LABEL: define void @narrow_store_user_mask_operand(
+; CHECK-SAME: i32 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP_PH:.*]]
+; CHECK:       [[LOOP_PH]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP_TAIL:.*]] ]
+; CHECK-NEXT:    [[X_POS:%.*]] = icmp sgt i32 [[X]], 0
+; CHECK-NEXT:    br i1 [[X_POS]], label %[[LOOP_BODY:.*]], label %[[LOOP_TAIL]]
+; CHECK:       [[LOOP_BODY]]:
+; CHECK-NEXT:    [[LD_P:%.*]] = load double, ptr @p, align 8
+; CHECK-NEXT:    [[GEP_Q_IV:%.*]] = getelementptr double, ptr @q, i64 [[IV]]
+; CHECK-NEXT:    [[GEP_Q_IV_8:%.*]] = getelementptr i8, ptr [[GEP_Q_IV]], i64 -8
+; CHECK-NEXT:    store double [[LD_P]], ptr [[GEP_Q_IV_8]], align 8
+; CHECK-NEXT:    br label %[[LOOP_TAIL]]
+; CHECK:       [[LOOP_TAIL]]:
+; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i64 [[IV]], 1
+; CHECK-NEXT:    br i1 [[EC]], label %[[EXIT:.*]], label %[[LOOP_PH]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop.ph
+
+loop.ph:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop.tail ]
+  %x.pos = icmp sgt i32 %x, 0
+  br i1 %x.pos, label %loop.body, label %loop.tail
+
+loop.body:
+  %ld.p = load double, ptr @p
+  %gep.q.iv = getelementptr double, ptr @q, i64 %iv
+  %gep.q.iv.8 = getelementptr i8, ptr %gep.q.iv, i64 -8
+  store double %ld.p, ptr %gep.q.iv.8
+  br label %loop.tail
+
+loop.tail:
+  %iv.next = add i64 %iv, 1
+  %ec = icmp eq i64 %iv, 1
+  br i1 %ec, label %exit, label %loop.ph
+
+exit:
+  ret void
+}