[VectorCombine] Don't shrink lshr if the shamt is not less than bitwidth

Author: dtcxzyw

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;
 

llvm/test/Transforms/VectorCombine/AArch64/shrink-types.ll

@@ -103,9 +103,8 @@ vector.body:
 define <2 x i32> @pr108698(<2 x i64> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @pr108698(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i64> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc <2 x i32> [[Y:%.*]] to <2 x i1>
-; CHECK-NEXT:    [[TMP2:%.*]] = lshr <2 x i1> [[CMP]], [[TMP1]]
-; CHECK-NEXT:    [[LSHR:%.*]] = zext <2 x i1> [[TMP2]] to <2 x i32>
+; CHECK-NEXT:    [[EXT:%.*]] = zext <2 x i1> [[CMP]] to <2 x i32>
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[EXT]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[LSHR]]
 ;
   %cmp = icmp eq <2 x i64> %x, zeroinitializer