[RISCV] Fold (X & -(1 << C1) & 0xffffffff) == C2 << C1 to sraiw X, C1 == C2.

[topperc]

We had an existing fold for (X & -(1 << C1) & 0xffffffff) == 0
which we can generalize to support comparing to constants other
than 0.

Previously we used srliw, but this generalizes better using sraiw.
I'm restricting to the case where C2 is simm12 or 2048 to allow
sraiw+addi/xori+seqz/snez to be used. Other constants require a
more careful analysis of the constants involved.

[llvmbot]

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

Author: Craig Topper (topperc)

Changes

We had an existing fold for (X & -(1 << C1) & 0xffffffff) == 0
which we can generalize to support comparing to constants other
than 0.

Previously we used srliw, but this generalizes better using sraiw.
I'm restricting to the case where C2 is simm12 or 2048 to allow
sraiw+addi/xori+seqz/snez to be used. Other constants require a
more careful analysis of the constants involved.

---

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

3 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+18-11)
• (modified) llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll (+59-2)
• (modified) llvm/test/CodeGen/RISCV/bittest.ll (+3-3)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 46f544c0d4df5..b5f2415fba5a5 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -16761,29 +16761,36 @@ static SDValue performSETCCCombine(SDNode *N,
           combineVectorSizedSetCCEquality(VT, N0, N1, Cond, dl, DAG, Subtarget))
     return V;
 
-  if (DCI.isAfterLegalizeDAG() && isNullConstant(N1) &&
+  if (DCI.isAfterLegalizeDAG() && isa<ConstantSDNode>(N1) &&
       N0.getOpcode() == ISD::AND && N0.hasOneUse() &&
       isa<ConstantSDNode>(N0.getOperand(1))) {
-    const APInt &AndRHSC =
-        cast<ConstantSDNode>(N0.getOperand(1))->getAPIntValue();
+    const APInt &AndRHSC = N0.getConstantOperandAPInt(1);
     // (X & -(1 << C)) == 0 -> (X >> C) == 0 if the AND constant can't use ANDI.
-    if (!isInt<12>(AndRHSC.getSExtValue()) && AndRHSC.isNegatedPowerOf2()) {
+    if (isNullConstant(N1) && !isInt<12>(AndRHSC.getSExtValue()) &&
+        AndRHSC.isNegatedPowerOf2()) {
       unsigned ShiftBits = AndRHSC.countr_zero();
       SDValue Shift = DAG.getNode(ISD::SRL, dl, OpVT, N0.getOperand(0),
                                   DAG.getConstant(ShiftBits, dl, OpVT));
       return DAG.getSetCC(dl, VT, Shift, N1, Cond);
     }
 
-    // Similar to above but handling the lower 32 bits by using srliw.
-    // FIXME: Handle the case where N1 is non-zero.
+    // Similar to above but handling the lower 32 bits by using sraiw. Allow
+    // comparing with constants other than 0 if the constant can be folded into
+    // addi or xori after shifting.
     if (OpVT == MVT::i64 && AndRHSC.getZExtValue() <= 0xffffffff &&
         isPowerOf2_32(-uint32_t(AndRHSC.getZExtValue()))) {
       unsigned ShiftBits = llvm::countr_zero(AndRHSC.getZExtValue());
-      SDValue And = DAG.getNode(ISD::AND, dl, OpVT, N0.getOperand(0),
-                                DAG.getConstant(0xffffffff, dl, OpVT));
-      SDValue Shift = DAG.getNode(ISD::SRL, dl, OpVT, And,
-                                  DAG.getConstant(ShiftBits, dl, OpVT));
-      return DAG.getSetCC(dl, VT, Shift, N1, Cond);
+      int64_t NewC = SignExtend64<32>(cast<ConstantSDNode>(N1)->getSExtValue());
+      NewC >>= ShiftBits;
+      if (isInt<12>(NewC) || NewC == 2048) {
+        SDValue SExt =
+            DAG.getNode(ISD::SIGN_EXTEND_INREG, dl, OpVT, N0.getOperand(0),
+                        DAG.getValueType(MVT::i32));
+        SDValue Shift = DAG.getNode(ISD::SRA, dl, OpVT, SExt,
+                                    DAG.getConstant(ShiftBits, dl, OpVT));
+        return DAG.getSetCC(dl, VT, Shift,
+                            DAG.getSignedConstant(NewC, dl, OpVT), Cond);
+      }
     }
   }
 
diff --git a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
index c39aedc341acc..55b0d1f0bf7be 100644
--- a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
+++ b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
@@ -74,7 +74,7 @@ define i1 @test5(i64 %x) {
 ;
 ; RV64-LABEL: test5:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    srliw a0, a0, 29
+; RV64-NEXT:    sraiw a0, a0, 29
 ; RV64-NEXT:    seqz a0, a0
 ; RV64-NEXT:    ret
   %a = and i64 %x, u0xE0000000
@@ -91,10 +91,67 @@ define i1 @test6(i64 %x) {
 ;
 ; RV64-LABEL: test6:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    srliw a0, a0, 29
+; RV64-NEXT:    sraiw a0, a0, 29
 ; RV64-NEXT:    snez a0, a0
 ; RV64-NEXT:    ret
   %a = and i64 %x, u0xE0000000
   %b = icmp ne i64 %a, 0
   ret i1 %b
 }
+
+define i1 @test7(i64 %x) {
+; RV32-LABEL: test7:
+; RV32:       # %bb.0:
+; RV32-NEXT:    srli a0, a0, 29
+; RV32-NEXT:    addi a0, a0, -6
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test7:
+; RV64:       # %bb.0:
+; RV64-NEXT:    sraiw a0, a0, 29
+; RV64-NEXT:    addi a0, a0, 2
+; RV64-NEXT:    seqz a0, a0
+; RV64-NEXT:    ret
+  %a = and i64 %x, u0xE0000000
+  %b = icmp eq i64 %a, u0xC0000000
+  ret i1 %b
+}
+
+define i1 @test8(i64 %x) {
+; RV32-LABEL: test8:
+; RV32:       # %bb.0:
+; RV32-NEXT:    srai a0, a0, 20
+; RV32-NEXT:    xori a0, a0, -2048
+; RV32-NEXT:    snez a0, a0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test8:
+; RV64:       # %bb.0:
+; RV64-NEXT:    sraiw a0, a0, 20
+; RV64-NEXT:    xori a0, a0, -2048
+; RV64-NEXT:    snez a0, a0
+; RV64-NEXT:    ret
+  %a = and i64 %x, u0xFFF00000
+  %b = icmp ne i64 %a, u0x80000000
+  ret i1 %b
+}
+
+define i1 @test9(i64 %x) {
+; RV32-LABEL: test9:
+; RV32:       # %bb.0:
+; RV32-NEXT:    srli a0, a0, 16
+; RV32-NEXT:    addi a0, a0, -2048
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test9:
+; RV64:       # %bb.0:
+; RV64-NEXT:    sraiw a0, a0, 16
+; RV64-NEXT:    addi a0, a0, -2048
+; RV64-NEXT:    seqz a0, a0
+; RV64-NEXT:    ret
+  %a = and i64 %x, u0xFFFF0000
+  %b = icmp eq i64 %a, u0x08000000
+  ret i1 %b
+}
diff --git a/llvm/test/CodeGen/RISCV/bittest.ll b/llvm/test/CodeGen/RISCV/bittest.ll
index 5d40d8f2d860d..35d38524c2e9a 100644
--- a/llvm/test/CodeGen/RISCV/bittest.ll
+++ b/llvm/test/CodeGen/RISCV/bittest.ll
@@ -187,13 +187,13 @@ define i64 @bittest_31_i64(i64 %a) nounwind {
 ;
 ; RV64ZBS-LABEL: bittest_31_i64:
 ; RV64ZBS:       # %bb.0:
-; RV64ZBS-NEXT:    srliw a0, a0, 31
+; RV64ZBS-NEXT:    sraiw a0, a0, 31
 ; RV64ZBS-NEXT:    seqz a0, a0
 ; RV64ZBS-NEXT:    ret
 ;
 ; RV64XTHEADBS-LABEL: bittest_31_i64:
 ; RV64XTHEADBS:       # %bb.0:
-; RV64XTHEADBS-NEXT:    srliw a0, a0, 31
+; RV64XTHEADBS-NEXT:    sraiw a0, a0, 31
 ; RV64XTHEADBS-NEXT:    seqz a0, a0
 ; RV64XTHEADBS-NEXT:    ret
   %shr = lshr i64 %a, 31
@@ -3517,7 +3517,7 @@ define i32 @bittest_31_andeq0_i64(i64 %x) {
 ;
 ; RV64-LABEL: bittest_31_andeq0_i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    srliw a0, a0, 31
+; RV64-NEXT:    sraiw a0, a0, 31
 ; RV64-NEXT:    seqz a0, a0
 ; RV64-NEXT:    ret
   %and = and i64 %x, 2147483648

[wangpc-pp]

LGTM.

[pfusik]

Here we discard bits 32-63 of N1. Don't we need to check them?

[topperc]

We also discarded the shifted out bits, we need to check those too. N1 needs to be a subset of AndRHSC. Thanks for catching that.

[pfusik]

LGTM