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

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