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

[topperc]

If using srl 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.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-risc-v

Author: Craig Topper (topperc)

Changes

If using srl 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.

---

Full diff: https://github.com/llvm/llvm-project/pull/151762.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp (+16-1)
• (added) llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll (+14)

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
new file mode 100644
index 0000000000000..be3de37927564
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
@@ -0,0 +1,14 @@
+; 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:    srai a0, a0, 30
+; CHECK-NEXT:    addi a0, a0, 2
+; CHECK-NEXT:    seqz a0, a0
+; CHECK-NEXT:    ret
+  %a = and i64 %x, -1073741824
+  %b = icmp eq i64 %a, -2147483648
+  ret i1 %b
+}

[RKSimon]

LGTM