[llvm] 4baf29e - [DAG] Handle cases where a shift amount is larger than the pre-extended value bitwidth

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-08-27T18:12:24+01:00
New Revision: 4baf29e81e30e6ebc1da56b9e5c338014961acf9

URL: https://github.com/llvm/llvm-project/commit/4baf29e81e30e6ebc1da56b9e5c338014961acf9
DIFF: https://github.com/llvm/llvm-project/commit/4baf29e81e30e6ebc1da56b9e5c338014961acf9.diff

LOG: [DAG] Handle cases where a shift amount is larger than the pre-extended value bitwidth

In the (zext (shl (zext x), cst)) -> (shl (zext x), cst) fold, don't use a bitmask / MaskedValueIsZero as we can't guarantee that the shift amount is in bounds.

Fixes #106202

Added: 
    llvm/test/CodeGen/SystemZ/pr106202.ll

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index b27f06f94ff0e7..b0a906743f29ff 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -14114,13 +14114,10 @@ SDValue DAGCombiner::visitZERO_EXTEND(SDNode *N) {
           if (ShAmtC->getAPIntValue().ugt(KnownZeroBits)) {
             // If the shift is too large, then see if we can deduce that the
             // shift is safe anyway.
-            // Create a mask that has ones for the bits being shifted out.
-            APInt ShiftOutMask =
-                APInt::getHighBitsSet(ShVal.getValueSizeInBits(),
-                                      ShAmtC->getAPIntValue().getZExtValue());
 
             // Check if the bits being shifted out are known to be zero.
-            if (!DAG.MaskedValueIsZero(ShVal, ShiftOutMask))
+            KnownBits KnownShVal = DAG.computeKnownBits(ShVal);
+            if (ShAmtC->getAPIntValue().ugt(KnownShVal.countMinLeadingZeros()))
               return SDValue();
           }
         }

diff  --git a/llvm/test/CodeGen/SystemZ/pr106202.ll b/llvm/test/CodeGen/SystemZ/pr106202.ll
new file mode 100644
index 00000000000000..7ec165e932eace
--- /dev/null
+++ b/llvm/test/CodeGen/SystemZ/pr106202.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=s390x-ibm-linux -mcpu=z10 | FileCheck %s
+
+ at g_0 = external dso_local local_unnamed_addr global i16, align 2
+ at g_1 = external dso_local local_unnamed_addr global i32, align 4
+ at g_2 = external dso_local local_unnamed_addr global i32, align 4
+
+define void @func() {
+; CHECK-LABEL: func:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lhi %r0, 41
+; CHECK-NEXT:    strl %r0, g_1
+; CHECK-NEXT:    lhi %r0, 0
+; CHECK-NEXT:    strl %r0, g_2
+; CHECK-NEXT:    br %r14
+  store i32 41, ptr @g_1, align 4
+  %1 = load i32, ptr @g_1, align 4
+  %2 = load i16, ptr @g_0, align 2
+  %3 = zext i16 %2 to i32
+  %4 = shl i32 %3, %1
+  %5 = zext i32 %4 to i64
+  %6 = shl i64 %5, 48
+  %7 = ashr exact i64 %6, 48
+  %8 = or i64 %7, 0
+  %9 = sext i32 %1 to i64
+  %10 = icmp sge i64 %8, %9
+  %11 = zext i1 %10 to i32
+  %12 = or i32 0, %11
+  store i32 %12, ptr @g_2, align 4
+  ret void
+}