[llvm] [TargetLowering][RISCV] Use sra for (X & -256) == 256 -> (X >> 8) == 1 if it yields a better icmp constant. (PR #151762)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/151762

If using lsr does not produce a legal constant for the RHS of the
final compare, try to use sra instead.
    
Because the AND constant is negative, the sign bits participate in the
compare. Using an arithmetic shift right duplicates that bit.

>From 3c8d224d47d602e9ac366ab2a7cd8387b66758ad Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Fri, 1 Aug 2025 12:42:28 -0700
Subject: [PATCH 1/2] Pre-commit test

---
 llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll

diff --git a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
new file mode 100644
index 0000000000000..681515c8fa221
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=riscv64 | FileCheck %s
+
+define i1 @src(i64 %x) {
+; CHECK-LABEL: src:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    srli a0, a0, 30
+; CHECK-NEXT:    lui a1, 524288
+; CHECK-NEXT:    srli a1, a1, 30
+; CHECK-NEXT:    xor a0, a0, a1
+; CHECK-NEXT:    seqz a0, a0
+; CHECK-NEXT:    ret
+  %a = and i64 %x, -1073741824
+  %b = icmp eq i64 %a, -2147483648
+  ret i1 %b
+}

>From 698ee2b6438ca06b8709b6f738fcc35993f9e46c Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Fri, 1 Aug 2025 12:49:32 -0700
Subject: [PATCH 2/2] [TargetLowering][RISCV] Use sra for (X & -256) == 256 ->
 (X >> 8) == 1 if it yields a better icmp constant.

If using an unsigned shift does produce a legal constant for the RHS of the
final compare, try to use sra instead.

Because the AND constant is negative, the sign bits participate in the
compare. Using an arithmetic shift right duplicates that bit.
---
 .../lib/CodeGen/SelectionDAG/TargetLowering.cpp | 17 ++++++++++++++++-
 llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll      |  6 ++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 48d6b9996e572..fbd4c1f4a0d85 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -5363,10 +5363,25 @@ SDValue TargetLowering::SimplifySetCC(EVT VT, SDValue N0, SDValue N1,
           if (AndRHSC.isNegatedPowerOf2() && C1.isSubsetOf(AndRHSC)) {
             unsigned ShiftBits = AndRHSC.countr_zero();
             if (!shouldAvoidTransformToShift(ShValTy, ShiftBits)) {
+              // If using an unsigned shift doesn't yield a legal compare
+              // immediate, try using sra instead.
+              APInt NewC = C1.lshr(ShiftBits);
+              if (NewC.getSignificantBits() > 64 ||
+                  !isLegalICmpImmediate(NewC.getSExtValue())) {
+                APInt SignedC = C1.ashr(ShiftBits);
+                if (SignedC.getSignificantBits() <= 64 &&
+                    isLegalICmpImmediate(SignedC.getSExtValue())) {
+                  SDValue Shift = DAG.getNode(
+                      ISD::SRA, dl, ShValTy, N0.getOperand(0),
+                      DAG.getShiftAmountConstant(ShiftBits, ShValTy, dl));
+                  SDValue CmpRHS = DAG.getConstant(SignedC, dl, ShValTy);
+                  return DAG.getSetCC(dl, VT, Shift, CmpRHS, Cond);
+                }
+              }
               SDValue Shift = DAG.getNode(
                   ISD::SRL, dl, ShValTy, N0.getOperand(0),
                   DAG.getShiftAmountConstant(ShiftBits, ShValTy, dl));
-              SDValue CmpRHS = DAG.getConstant(C1.lshr(ShiftBits), dl, ShValTy);
+              SDValue CmpRHS = DAG.getConstant(NewC, dl, ShValTy);
               return DAG.getSetCC(dl, VT, Shift, CmpRHS, Cond);
             }
           }
diff --git a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
index 681515c8fa221..be3de37927564 100644
--- a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
+++ b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
@@ -4,10 +4,8 @@
 define i1 @src(i64 %x) {
 ; CHECK-LABEL: src:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    srli a0, a0, 30
-; CHECK-NEXT:    lui a1, 524288
-; CHECK-NEXT:    srli a1, a1, 30
-; CHECK-NEXT:    xor a0, a0, a1
+; CHECK-NEXT:    srai a0, a0, 30
+; CHECK-NEXT:    addi a0, a0, 2
 ; CHECK-NEXT:    seqz a0, a0
 ; CHECK-NEXT:    ret
   %a = and i64 %x, -1073741824