[SLP] Add a check for poison value in AShrChecker

Need to check if the value in AShrChecker is a poison before casting it
to instruction to avoid compiler crash

Fixes #118030

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -17799,6 +17799,8 @@ bool BoUpSLP::collectValuesToDemote(
     // original type and the sign bit of the truncate type are similar.
     auto AShrChecker = [&](unsigned BitWidth, unsigned OrigBitWidth) {
       return all_of(E.Scalars, [&](Value *V) {
+        if (isa<PoisonValue>(V))
+          return true;
         auto *I = cast<Instruction>(V);
         KnownBits AmtKnownBits = computeKnownBits(I->getOperand(1), *DL);
         unsigned ShiftedBits = OrigBitWidth - BitWidth;

llvm/test/Transforms/SLPVectorizer/X86/ashr-node-with-poison.ll

@@ -0,0 +1,120 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S --passes=slp-vectorizer -mtriple=x86_64-unknown-linux < %s | FileCheck %s
+
+define i32 @test(ptr %n, i32 %conv57, i1 %tobool5.not, i64 %bf.load14) {
+; CHECK-LABEL: define i32 @test(
+; CHECK-SAME: ptr [[N:%.*]], i32 [[CONV57:%.*]], i1 [[TOBOOL5_NOT:%.*]], i64 [[BF_LOAD14:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br i1 false, label %[[ENTRY_IF_END54_CRIT_EDGE:.*]], label %[[WHILE_COND:.*]]
+; CHECK:       [[ENTRY_IF_END54_CRIT_EDGE]]:
+; CHECK-NEXT:    br label %[[IF_END54:.*]]
+; CHECK:       [[TTHREAD_PRE_SPLIT:.*]]:
+; CHECK-NEXT:    br label %[[T:.*]]
+; CHECK:       [[T]]:
+; CHECK-NEXT:    [[W_2:%.*]] = phi i32 [ 1, %[[TTHREAD_PRE_SPLIT]] ], [ 0, %[[IF_END83:.*]] ]
+; CHECK-NEXT:    br label %[[IF_END7:.*]]
+; CHECK:       [[T_U_CRIT_EDGE:.*]]:
+; CHECK-NEXT:    br label %[[IF_END7]]
+; CHECK:       [[IF_END7]]:
+; CHECK-NEXT:    [[W_4:%.*]] = phi i32 [ [[W_2]], %[[T]] ], [ [[W_2]], %[[T_U_CRIT_EDGE]] ], [ [[BF_CAST25:%.*]], %[[WHILE_BODY:.*]] ]
+; CHECK-NEXT:    [[A_4:%.*]] = phi i32 [ 0, %[[T]] ], [ [[CONV57]], %[[T_U_CRIT_EDGE]] ], [ [[BF_CAST2910:%.*]], %[[WHILE_BODY]] ]
+; CHECK-NEXT:    [[B_4:%.*]] = phi i32 [ 0, %[[T]] ], [ 0, %[[T_U_CRIT_EDGE]] ], [ [[BF_CAST2910]], %[[WHILE_BODY]] ]
+; CHECK-NEXT:    [[C_4:%.*]] = phi i32 [ 0, %[[T]] ], [ 1, %[[T_U_CRIT_EDGE]] ], [ poison, %[[WHILE_BODY]] ]
+; CHECK-NEXT:    br label %[[V:.*]]
+; CHECK:       [[WHILE_COND]]:
+; CHECK-NEXT:    [[BF_LOAD66_PRE_PRE1135:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[SPEC_SELECT:%.*]], %[[IF_END42:.*]] ]
+; CHECK-NEXT:    br i1 [[TOBOOL5_NOT]], label %[[IF_END54]], label %[[WHILE_BODY]]
+; CHECK:       [[WHILE_BODY]]:
+; CHECK-NEXT:    [[BF_ASHR24:%.*]] = ashr i64 [[BF_LOAD14]], 33
+; CHECK-NEXT:    [[BF_CAST25]] = trunc nsw i64 [[BF_ASHR24]] to i32
+; CHECK-NEXT:    [[BF_ASHR28:%.*]] = lshr i64 [[BF_LOAD14]], 1
+; CHECK-NEXT:    [[BF_CAST2910]] = trunc i64 [[BF_ASHR28]] to i32
+; CHECK-NEXT:    br label %[[IF_END7]]
+; CHECK:       [[IF_END36:.*]]:
+; CHECK-NEXT:    br label %[[V]]
+; CHECK:       [[V]]:
+; CHECK-NEXT:    [[C_7:%.*]] = phi i32 [ [[C_4]], %[[IF_END7]] ], [ 0, %[[IF_END36]] ]
+; CHECK-NEXT:    br i1 true, label %[[IF_END42]], label %[[V_IF_END83_CRIT_EDGE:.*]]
+; CHECK:       [[V_IF_END83_CRIT_EDGE]]:
+; CHECK-NEXT:    br label %[[IF_END83]]
+; CHECK:       [[IF_END42]]:
+; CHECK-NEXT:    [[TOBOOL43_NOT:%.*]] = icmp eq i32 [[B_4]], 0
+; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[TOBOOL43_NOT]], i32 0, i32 [[W_4]]
+; CHECK-NEXT:    [[SPEC_SELECT]] = zext i32 [[NARROW]] to i64
+; CHECK-NEXT:    [[BF_VALUE48:%.*]] = zext i32 [[A_4]] to i64
+; CHECK-NEXT:    store i64 [[BF_VALUE48]], ptr [[N]], align 8
+; CHECK-NEXT:    store i32 [[C_7]], ptr [[N]], align 4
+; CHECK-NEXT:    br label %[[WHILE_COND]]
+; CHECK:       [[IF_END54]]:
+; CHECK-NEXT:    [[BF_LOAD66_PRE_PRE113125:%.*]] = phi i64 [ [[BF_LOAD66_PRE_PRE1135]], %[[WHILE_COND]] ], [ poison, %[[ENTRY_IF_END54_CRIT_EDGE]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 [[BF_LOAD66_PRE_PRE113125]], 0
+; CHECK-NEXT:    br i1 [[TMP0]], label %[[IF_END83]], label %[[AI_IF_END76_CRIT_EDGE:.*]]
+; CHECK:       [[AI_IF_END76_CRIT_EDGE]]:
+; CHECK-NEXT:    br label %[[IF_END83]]
+; CHECK:       [[IF_END83]]:
+; CHECK-NEXT:    br label %[[T]]
+;
+entry:
+  br i1 false, label %entry.if.end54_crit_edge, label %while.cond
+
+entry.if.end54_crit_edge:
+  br label %if.end54
+
+tthread-pre-split:
+  br label %t
+
+t:
+  %w.2 = phi i32 [ 1, %tthread-pre-split ], [ 0, %if.end83 ]
+  br label %if.end7
+
+t.u_crit_edge:
+  br label %if.end7
+
+if.end7:
+  %w.4 = phi i32 [ %w.2, %t ], [ %w.2, %t.u_crit_edge ], [ %bf.cast25, %while.body ]
+  %a.4 = phi i32 [ 0, %t ], [ %conv57, %t.u_crit_edge ], [ %bf.cast2910, %while.body ]
+  %b.4 = phi i32 [ 0, %t ], [ 0, %t.u_crit_edge ], [ %bf.cast2910, %while.body ]
+  %c.4 = phi i32 [ 0, %t ], [ 1, %t.u_crit_edge ], [ poison, %while.body ]
+  br label %v
+
+while.cond:
+  %bf.load66.pre.pre1135 = phi i64 [ 0, %entry ], [ %spec.select, %if.end42 ]
+  br i1 %tobool5.not, label %if.end54, label %while.body
+
+while.body:
+  %bf.ashr24 = ashr i64 %bf.load14, 33
+  %bf.cast25 = trunc nsw i64 %bf.ashr24 to i32
+  %bf.ashr28 = lshr i64 %bf.load14, 1
+  %bf.cast2910 = trunc i64 %bf.ashr28 to i32
+  br label %if.end7
+
+if.end36:
+  br label %v
+
+v:
+  %c.7 = phi i32 [ %c.4, %if.end7 ], [ 0, %if.end36 ]
+  br i1 true, label %if.end42, label %v.if.end83_crit_edge
+
+v.if.end83_crit_edge:
+  br label %if.end83
+
+if.end42:
+  %tobool43.not = icmp eq i32 %b.4, 0
+  %narrow = select i1 %tobool43.not, i32 0, i32 %w.4
+  %spec.select = zext i32 %narrow to i64
+  %bf.value48 = zext i32 %a.4 to i64
+  store i64 %bf.value48, ptr %n, align 8
+  store i32 %c.7, ptr %n, align 4
+  br label %while.cond
+
+if.end54:
+  %bf.load66.pre.pre113125 = phi i64 [ %bf.load66.pre.pre1135, %while.cond ], [ poison, %entry.if.end54_crit_edge ]
+  %0 = icmp eq i64 %bf.load66.pre.pre113125, 0
+  br i1 %0, label %if.end83, label %ai.if.end76_crit_edge
+
+ai.if.end76_crit_edge:
+  br label %if.end83
+
+if.end83:
+  br label %t
+}