[llvm] Port isNonZeroShift to SelectionDAG (PR #146125)

via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 27 13:12:37 PDT 2025


https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/146125

>From 838f11e706aa5c0aa2f724b65877f1056ab950d7 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Fri, 27 Jun 2025 15:45:17 -0400
Subject: [PATCH 1/3] Pre-commit test (NFC)

---
 llvm/test/CodeGen/AArch64/arm64-fold-lshr.ll | 53 ++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/llvm/test/CodeGen/AArch64/arm64-fold-lshr.ll b/llvm/test/CodeGen/AArch64/arm64-fold-lshr.ll
index 9dfc8df703ce6..0dd868f1c7da9 100644
--- a/llvm/test/CodeGen/AArch64/arm64-fold-lshr.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-fold-lshr.ll
@@ -136,3 +136,56 @@ entry:
   %0 = load i64, ptr %arrayidx, align 8
   ret i64 %0
 }
+
+define i32 @ctpop_shl2_1_nz(i32 %x) {
+; CHECK-LABEL: ctpop_shl2_1_nz:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #2 // =0x2
+; CHECK-NEXT:    and w9, w0, #0xf
+; CHECK-NEXT:    lsl w8, w8, w9
+; CHECK-NEXT:    fmov s0, w8
+; CHECK-NEXT:    cnt v0.8b, v0.8b
+; CHECK-NEXT:    addv b0, v0.8b
+; CHECK-NEXT:    fmov w0, s0
+; CHECK-NEXT:    ret
+  %and = and i32 %x, 15
+  %shl = shl i32 2, %and
+  %cnt = call i32 @llvm.ctpop.i32(i32 %shl)
+  ret i32 %cnt
+}
+
+define i1 @ashr_nz_bounded_cnt (i32 %x, i32 %y) {
+; CHECK-LABEL: ashr_nz_bounded_cnt:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #16 // =0x10
+; CHECK-NEXT:    and w9, w0, #0x10
+; CHECK-NEXT:    movk w8, #6146, lsl #16
+; CHECK-NEXT:    orr w8, w1, w8
+; CHECK-NEXT:    asr w8, w8, w9
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    cset w0, eq
+; CHECK-NEXT:    ret
+  %cnt = and i32 %x, 16
+  %val = or i32 %y, 402784272
+  %shl = ashr i32 %val, %cnt
+  %r = icmp eq i32 %shl, 0
+  ret i1 %r
+}
+
+define i1 @ashr_nz_bounded_cnt_vec(i32 %x, i32 %y) {
+; CHECK-LABEL: ashr_nz_bounded_cnt_vec:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #16 // =0x10
+; CHECK-NEXT:    and w9, w0, #0x10
+; CHECK-NEXT:    movk w8, #6146, lsl #16
+; CHECK-NEXT:    orr w8, w1, w8
+; CHECK-NEXT:    lsr w8, w8, w9
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    cset w0, eq
+; CHECK-NEXT:    ret
+  %cnt = and i32 %x, 16
+  %val = or i32 %y, 402784272
+  %shl = lshr i32 %val, %cnt
+  %r = icmp eq i32 %shl, 0
+  ret i1 %r
+}

>From acd384674448525c3d8cc7df43f82c3f52135e79 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Fri, 27 Jun 2025 13:58:29 -0400
Subject: [PATCH 2/3] Port isNonZeroShift to SelectionDAG

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 59 ++++++++++++++++---
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index fe9a6ea3e77e6..c4471ed42a0d4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5871,6 +5871,55 @@ bool SelectionDAG::isKnownNeverZeroFloat(SDValue Op) const {
       Op, [](ConstantFPSDNode *C) { return !C->isZero(); });
 }
 
+static bool isNonZeroShift(const SelectionDAG &DAG, const SDValue I,
+                           const KnownBits &KnownVal, unsigned Depth) {
+  auto ShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I.getOpcode()) {
+    case Instruction::Shl:
+      return Lhs.shl(Rhs);
+    case Instruction::LShr:
+      return Lhs.lshr(Rhs);
+    case Instruction::AShr:
+      return Lhs.ashr(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  auto InvShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
+    switch (I.getOpcode()) {
+    case ISD::SHL:
+      return Lhs.lshr(Rhs);
+    case ISD::SRA:
+    case ISD::SRL:
+      return Lhs.shl(Rhs);
+    default:
+      llvm_unreachable("Unknown Shift Opcode");
+    }
+  };
+
+  if (KnownVal.isUnknown())
+    return false;
+
+  KnownBits KnownCnt = DAG.computeKnownBits(I.getOperand(1), Depth + 1);
+  APInt MaxShift = KnownCnt.getMaxValue();
+  unsigned NumBits = KnownVal.getBitWidth();
+  if (MaxShift.uge(NumBits))
+    return false;
+
+  if (!ShiftOp(KnownVal.One, MaxShift).isZero())
+    return true;
+
+  // If all of the bits shifted out are known to be zero, and Val is known
+  // non-zero then at least one non-zero bit must remain.
+  if (InvShiftOp(KnownVal.Zero, NumBits - MaxShift)
+          .eq(InvShiftOp(APInt::getAllOnes(NumBits), NumBits - MaxShift)) &&
+      DAG.isKnownNeverZero(I.getOperand(0), Depth + 1))
+    return true;
+
+  return false;
+}
+
 bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
   if (Depth >= MaxRecursionDepth)
     return false; // Limit search depth.
@@ -5906,9 +5955,7 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     if (ValKnown.One[0])
       return true;
     // If max shift cnt of known ones is non-zero, result is non-zero.
-    APInt MaxCnt = computeKnownBits(Op.getOperand(1), Depth + 1).getMaxValue();
-    if (MaxCnt.ult(ValKnown.getBitWidth()) &&
-        !ValKnown.One.shl(MaxCnt).isZero())
+    if (isNonZeroShift(*this, Op, ValKnown, Depth))
       return true;
     break;
   }
@@ -5968,10 +6015,8 @@ bool SelectionDAG::isKnownNeverZero(SDValue Op, unsigned Depth) const {
     KnownBits ValKnown = computeKnownBits(Op.getOperand(0), Depth + 1);
     if (ValKnown.isNegative())
       return true;
-    // If max shift cnt of known ones is non-zero, result is non-zero.
-    APInt MaxCnt = computeKnownBits(Op.getOperand(1), Depth + 1).getMaxValue();
-    if (MaxCnt.ult(ValKnown.getBitWidth()) &&
-        !ValKnown.One.lshr(MaxCnt).isZero())
+
+    if (isNonZeroShift(*this, Op, ValKnown, Depth))
       return true;
     break;
   }

>From be797b487f42bfc420796d01842db98001a584f4 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Fri, 27 Jun 2025 14:54:37 -0400
Subject: [PATCH 3/3] Fix

---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index c4471ed42a0d4..3074cc825e07a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5875,11 +5875,11 @@ static bool isNonZeroShift(const SelectionDAG &DAG, const SDValue I,
                            const KnownBits &KnownVal, unsigned Depth) {
   auto ShiftOp = [&](const APInt &Lhs, const APInt &Rhs) {
     switch (I.getOpcode()) {
-    case Instruction::Shl:
+    case ISD::SHL:
       return Lhs.shl(Rhs);
-    case Instruction::LShr:
+    case ISD::SRL:
       return Lhs.lshr(Rhs);
-    case Instruction::AShr:
+    case ISD::SRA:
       return Lhs.ashr(Rhs);
     default:
       llvm_unreachable("Unknown Shift Opcode");



More information about the llvm-commits mailing list