Bug type: miscompilation
Component: VectorCombine
Bug location: llvm/lib/Transforms/Vectorize/VectorCombine.cpp: VectorCombine::shrinkType
Issue title: wrong code at -O{2,3} on x86_64-linux-gnu
Labels: miscompilation
vectorizers
Reported by: zhendongsu

It appears to be a recent regression as it doesn't reproduce with 18.1.0 and earlier. 

Compiler Explorer: https://godbolt.org/z/EGasYh8x1

```
[691] % clangtk -v
clang version 20.0.0git (https://github.com/llvm/llvm-project.git 223e2efa5e886502a9467b7ef700ebce9b7886e8)
Target: x86_64-unknown-linux-gnu
Thread model: posix
InstalledDir: /local/suz-local/software/local/clang-trunk/bin
Build config: +assertions
Found candidate GCC installation: /usr/lib/gcc/x86_64-linux-gnu/10
Found candidate GCC installation: /usr/lib/gcc/x86_64-linux-gnu/11
Found candidate GCC installation: /usr/lib/gcc/x86_64-linux-gnu/9
Selected GCC installation: /usr/lib/gcc/x86_64-linux-gnu/11
Candidate multilib: .;@m64
Selected multilib: .;@m64
[692] % 
[692] % clangtk -O1 small.c; ./a.out
[693] % 
[693] % clangtk -O2 small.c
[694] % ./a.out
Aborted
[695] % 
[695] % cat small.c
char a[8];
int b = 1, c;
int main() {
  for (; c < 8; c++)
    a[c] = !c >> b;
  if (a[0] != 0)
    __builtin_abort();
  return 0;
}
```

Comments: author: dtcxzyw
body: Reproducer: https://alive2.llvm.org/ce/z/5SkPZS
```
; bin/opt -passes=vector-combine test.ll -S

target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

define <2 x i32> @test(<2 x i64> %vec.ind16, <2 x i32> %broadcast.splat20) {
  %19 = icmp eq <2 x i64> %vec.ind16, zeroinitializer
  %20 = zext <2 x i1> %19 to <2 x i32>
  %21 = lshr <2 x i32> %20, %broadcast.splat20
  ret <2 x i32> %21
}
```
```
define <2 x i32> @test(<2 x i64> %vec.ind16, <2 x i32> %broadcast.splat20) {
  %1 = icmp eq <2 x i64> %vec.ind16, zeroinitializer
  %2 = trunc <2 x i32> %broadcast.splat20 to <2 x i1>
  %3 = lshr <2 x i1> %1, %2
  %4 = zext <2 x i1> %3 to <2 x i32>
  ret <2 x i32> %4
}
```
author: dtcxzyw
body: Related PR: https://github.com/llvm/llvm-project/pull/104606


Fix patch:
commit 87663fdab9d0e7bcc0b963ea078da9e2eb574908
Author: Yingwei Zheng <dtcxzyw2333@gmail.com>
Date:   Sun Sep 15 18:38:06 2024 +0800

     [VectorCombine] Don't shrink lshr if the shamt is not less than bitwidth (#108705)
    
    Consider the following case:
    ```
    define <2 x i32> @test(<2 x i64> %vec.ind16, <2 x i32> %broadcast.splat20) {
      %19 = icmp eq <2 x i64> %vec.ind16, zeroinitializer
      %20 = zext <2 x i1> %19 to <2 x i32>
      %21 = lshr <2 x i32> %20, %broadcast.splat20
      ret <2 x i32> %21
    }
    ```
    After https://github.com/llvm/llvm-project/pull/104606, we shrink the
    lshr into:
    ```
    define <2 x i32> @test(<2 x i64> %vec.ind16, <2 x i32> %broadcast.splat20) {
      %1 = icmp eq <2 x i64> %vec.ind16, zeroinitializer
      %2 = trunc <2 x i32> %broadcast.splat20 to <2 x i1>
      %3 = lshr <2 x i1> %1, %2
      %4 = zext <2 x i1> %3 to <2 x i32>
      ret <2 x i32> %4
    }
    ```
    It is incorrect since `lshr i1 X, 1` returns `poison`.
    This patch adds additional check on the shamt operand. The lshr will get
    shrunk iff we ensure that the shamt is less than bitwidth of the smaller
    type. As `computeKnownBits(&I, *DL).countMaxActiveBits() > BW` always
    evaluates to true for `lshr(zext(X), Y)`, this check will only apply to
    bitwise logical instructions.
    
    Alive2: https://alive2.llvm.org/ce/z/j_RmTa
    Fixes https://github.com/llvm/llvm-project/issues/108698.

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index d7afe2f426d3..58701bfa60a3 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -2597,11 +2597,19 @@ bool VectorCombine::shrinkType(llvm::Instruction &I) {
   auto *SmallTy = cast<FixedVectorType>(ZExted->getType());
   unsigned BW = SmallTy->getElementType()->getPrimitiveSizeInBits();
 
-  // Check that the expression overall uses at most the same number of bits as
-  // ZExted
-  KnownBits KB = computeKnownBits(&I, *DL);
-  if (KB.countMaxActiveBits() > BW)
-    return false;
+  if (I.getOpcode() == Instruction::LShr) {
+    // Check that the shift amount is less than the number of bits in the
+    // smaller type. Otherwise, the smaller lshr will return a poison value.
+    KnownBits ShAmtKB = computeKnownBits(I.getOperand(1), *DL);
+    if (ShAmtKB.getMaxValue().uge(BW))
+      return false;
+  } else {
+    // Check that the expression overall uses at most the same number of bits as
+    // ZExted
+    KnownBits KB = computeKnownBits(&I, *DL);
+    if (KB.countMaxActiveBits() > BW)
+      return false;
+  }
 
   // Calculate costs of leaving current IR as it is and moving ZExt operation
   // later, along with adding truncates if needed
@@ -2628,7 +2636,7 @@ bool VectorCombine::shrinkType(llvm::Instruction &I) {
       return false;
 
     // Check if we can propagate ZExt through its other users
-    KB = computeKnownBits(UI, *DL);
+    KnownBits KB = computeKnownBits(UI, *DL);
     if (KB.countMaxActiveBits() > BW)
       return false;
 


Regression tests:
file: llvm/test/Transforms/VectorCombine/AArch64/shrink-types.ll
commands: opt -passes=vector-combine -S %s
tests: test_name: pr108698
test_body: target triple = "aarch64"

define <2 x i32> @pr108698(<2 x i64> %x, <2 x i32> %y) {
  %cmp = icmp eq <2 x i64> %x, zeroinitializer
  %ext = zext <2 x i1> %cmp to <2 x i32>
  %lshr = lshr <2 x i32> %ext, %y
  ret <2 x i32> %lshr
}