[llvm] [SLP]Fix narrowed leaf masks for demoted reduction root (PR #227009)

[Alexey Bataev via llvm-commits]

https://github.com/alexey-bataev created https://github.com/llvm/llvm-project/pull/227009

The root of the narrowed-leaf reduction may be demoted and resized to the
reduction bitwidth of the wide reduction operations, while the masks are
built for the narrow leaves type, causes the crash.

Fixes https://github.com/llvm/llvm-project/pull/224919#issuecomment-5869787812


>From 800baed9354c63b7de97e73b160a4d046cd8eb23 Mon Sep 17 00:00:00 2001
From: Alexey Bataev <a.bataev at outlook.com>
Date: Mon, 28 Sep 2026 07:51:07 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.7
---
 .../Transforms/Vectorize/SLPVectorizer.cpp    | 10 +++-
 .../X86/narrowed-reduction-minbw-root.ll      | 60 +++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/SLPVectorizer/X86/narrowed-reduction-minbw-root.ll

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index c9e649c551b92..198ed26bcc719 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -33074,8 +33074,16 @@ class HorizontalReduction {
         // the absorbed narrow shls are masked off before the widening.
         if (!NarrowedLeafShifts.empty()) {
           Type *WideTy = ReductionRoot->getType();
-          Type *NarrowTy = VectorizedRoot->getType()->getScalarType();
+          Type *NarrowTy = VL.front()->getType();
           unsigned VF = getNumElements(VectorizedRoot->getType());
+          // The root may be demoted and resized to the reduction bitwidth,
+          // which is computed for the wide reduction operations.
+          if (VectorizedRoot->getType()->getScalarType() != NarrowTy) {
+            VectorizedRoot = Builder.CreateIntCast(
+                VectorizedRoot, getWidenedType(NarrowTy, VF),
+                V.isSignedMinBitwidthRootNode());
+            ++NumVectorInstructions;
+          }
           SmallVector<Constant *> ShiftConsts(VF, ConstantInt::get(WideTy, 0));
           SmallVector<Constant *> MaskConsts(
               VF, Constant::getAllOnesValue(NarrowTy));
diff --git a/llvm/test/Transforms/SLPVectorizer/X86/narrowed-reduction-minbw-root.ll b/llvm/test/Transforms/SLPVectorizer/X86/narrowed-reduction-minbw-root.ll
new file mode 100644
index 0000000000000..38915fa3df345
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/narrowed-reduction-minbw-root.ll
@@ -0,0 +1,60 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=slp-vectorizer -S -mtriple=x86_64-unknown-linux-gnu -mcpu=x86-64-v3 < %s | FileCheck %s
+
+define i32 @test(i1 %c0, i1 %c1, i8 %a, i8 %b, i8 %d, i8 %e) {
+; CHECK-LABEL: define i32 @test(
+; CHECK-SAME: i1 [[C0:%.*]], i1 [[C1:%.*]], i8 [[A:%.*]], i8 [[B:%.*]], i8 [[D:%.*]], i8 [[E:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <2 x i1> poison, i1 [[C0]], i64 0
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i1> [[TMP0]], i1 [[C1]], i64 1
+; CHECK-NEXT:    br i1 [[C0]], label %[[EXIT:.*]], label %[[BB:.*]]
+; CHECK:       [[BB]]:
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[TMP2:%.*]] = phi <2 x i1> [ [[TMP1]], %[[ENTRY]] ], [ zeroinitializer, %[[BB]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = zext <2 x i1> [[TMP2]] to <2 x i16>
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc <2 x i16> [[TMP3]] to <2 x i8>
+; CHECK-NEXT:    [[TMP5:%.*]] = and <2 x i8> [[TMP4]], splat (i8 127)
+; CHECK-NEXT:    [[TMP6:%.*]] = zext <2 x i8> [[TMP5]] to <2 x i32>
+; CHECK-NEXT:    [[TMP7:%.*]] = shl <2 x i32> [[TMP6]], splat (i32 1)
+; CHECK-NEXT:    [[TMP8:%.*]] = call i32 @llvm.vector.reduce.or.v2i32(<2 x i32> [[TMP7]])
+; CHECK-NEXT:    [[TMP9:%.*]] = zext i8 [[B]] to i32
+; CHECK-NEXT:    [[TMP10:%.*]] = shl i32 [[TMP9]], 1
+; CHECK-NEXT:    [[OP_RDX:%.*]] = or i32 [[TMP8]], [[TMP10]]
+; CHECK-NEXT:    [[TMP11:%.*]] = and i8 [[D]], 1
+; CHECK-NEXT:    [[TMP12:%.*]] = zext i8 [[TMP11]] to i32
+; CHECK-NEXT:    [[TMP13:%.*]] = shl i32 [[TMP12]], 1
+; CHECK-NEXT:    [[TMP14:%.*]] = and i8 [[E]], 1
+; CHECK-NEXT:    [[TMP15:%.*]] = zext i8 [[TMP14]] to i32
+; CHECK-NEXT:    [[OP_RDX1:%.*]] = or i32 [[TMP13]], [[TMP15]]
+; CHECK-NEXT:    [[OP_RDX2:%.*]] = or i32 [[OP_RDX]], [[OP_RDX1]]
+; CHECK-NEXT:    [[TMP16:%.*]] = zext i8 [[A]] to i32
+; CHECK-NEXT:    [[OP_RDX3:%.*]] = or i32 [[OP_RDX2]], [[TMP16]]
+; CHECK-NEXT:    ret i32 [[OP_RDX3]]
+;
+entry:
+  %z0 = zext i1 %c0 to i8
+  %z1 = zext i1 %c1 to i8
+  br i1 %c0, label %exit, label %bb
+
+bb:
+  br label %exit
+
+exit:
+  %p0 = phi i8 [ %z0, %entry ], [ 0, %bb ]
+  %p1 = phi i8 [ %z1, %entry ], [ 0, %bb ]
+  %or0 = or i8 %p1, %p0
+  %dm = and i8 %d, 1
+  %or1 = or i8 %or0, %dm
+  %bz = zext i8 %b to i32
+  %bs = shl i32 %bz, 1
+  %sh = shl i8 %or1, 1
+  %shz = zext i8 %sh to i32
+  %r0 = or i32 %bs, %shz
+  %em = and i8 %e, 1
+  %emz = zext i8 %em to i32
+  %r1 = or i32 %r0, %emz
+  %az = zext i8 %a to i32
+  %r2 = or i32 %r1, %az
+  ret i32 %r2
+}