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

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.

Author: topperc

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -16761,29 +16761,37 @@ 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.
-    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);
+    // 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.
+    uint64_t N1Int = cast<ConstantSDNode>(N1)->getZExtValue();
+    uint64_t AndRHSInt = AndRHSC.getZExtValue();
+    if (OpVT == MVT::i64 && AndRHSInt <= 0xffffffff &&
+        isPowerOf2_32(-uint32_t(AndRHSInt)) && (N1Int & AndRHSInt) == N1Int) {
+      unsigned ShiftBits = llvm::countr_zero(AndRHSInt);
+      int64_t NewC = SignExtend64<32>(N1Int) >> ShiftBits;
+      if (NewC >= -2048 && 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);
+      }
     }
   }
 

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
+}

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