[RISCV] Return true in hasBitTest when Zbs is enabled and update BEXTI pattern for resulting canonicalisation

As the Zbs extension includes bext[i] for bit extract, we can
unconditionally return true from this hook. This hook causes the DAG
combiner to perform the following canonicalisation:

  and (not (srl X, C)), 1 --> (and X, 1<<C) == 0
  and (srl (not X), C)), 1 --> (and X, 1<<C) == 0

As simply changing the hook causes a codegen regression, this patch also
modifies a BEXTI pattern to match this canonicalised form.

As BSETINVMask is now used for BEXT as well as BSET and BINV, it has
been renamed to the more generic SingleBitSetMask.

There is one codegen change in bittest.ll for bittest_31_i64 (NOT+BEXTI
rather than NOT+SRLIW). This is neutral in terms of code quality.

Differential Revision: https://reviews.llvm.org/D131482

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -1178,6 +1178,9 @@ bool RISCVTargetLowering::hasAndNotCompare(SDValue Y) const {
 }
 
 bool RISCVTargetLowering::hasBitTest(SDValue X, SDValue Y) const {
+  // Zbs provides BEXT[_I], which can be used with SEQZ/SNEZ as a bit test.
+  if (Subtarget.hasStdExtZbs())
+    return X.getValueType().isScalarInteger();
   // We can use ANDI+SEQZ/SNEZ as a bit test. Y contains the bit position.
   auto *C = dyn_cast<ConstantSDNode>(Y);
   return C && C->getAPIntValue().ule(10);

llvm/lib/Target/RISCV/RISCVInstrInfoZb.td

@@ -89,7 +89,7 @@ def BCLRXForm : SDNodeXForm<imm, [{
                                    SDLoc(N), N->getValueType(0));
 }]>;
 
-def BSETINVXForm : SDNodeXForm<imm, [{
+def SingleBitSetMaskToIndex : SDNodeXForm<imm, [{
   // Find the lowest 1.
   return CurDAG->getTargetConstant(countTrailingZeros(N->getZExtValue()),
                                    SDLoc(N), N->getValueType(0));
@@ -103,11 +103,11 @@ def BCLRMask : ImmLeaf<XLenVT, [{
 }], BCLRXForm>;
 
 // Checks if this mask has a single 1 bit and cannot be used with ORI/XORI.
-def BSETINVMask : ImmLeaf<XLenVT, [{
+def SingleBitSetMask : ImmLeaf<XLenVT, [{
   if (Subtarget->is64Bit())
     return !isInt<12>(Imm) && isPowerOf2_64(Imm);
   return !isInt<12>(Imm) && isPowerOf2_32(Imm);
-}], BSETINVXForm>;
+}], SingleBitSetMaskToIndex>;
 
 // Check if (or r, i) can be optimized to (BSETI (BSETI r, i0), i1),
 // in which i = (1 << i0) | (1 << i1).
@@ -885,16 +885,16 @@ def : Pat<(shiftop<shl> 1, GPR:$rs2),
 
 def : Pat<(and GPR:$rs1, BCLRMask:$mask),
           (BCLRI GPR:$rs1, BCLRMask:$mask)>;
-def : Pat<(or GPR:$rs1, BSETINVMask:$mask),
-          (BSETI GPR:$rs1, BSETINVMask:$mask)>;
-def : Pat<(xor GPR:$rs1, BSETINVMask:$mask),
-          (BINVI GPR:$rs1, BSETINVMask:$mask)>;
+def : Pat<(or GPR:$rs1, SingleBitSetMask:$mask),
+          (BSETI GPR:$rs1, SingleBitSetMask:$mask)>;
+def : Pat<(xor GPR:$rs1, SingleBitSetMask:$mask),
+          (BINVI GPR:$rs1, SingleBitSetMask:$mask)>;
 
 def : Pat<(and (srl GPR:$rs1, uimmlog2xlen:$shamt), (XLenVT 1)),
           (BEXTI GPR:$rs1, uimmlog2xlen:$shamt)>;
 
-def : Pat<(and (not (srl GPR:$rs1, uimmlog2xlen:$shamt)), (XLenVT 1)),
-          (XORI (BEXTI GPR:$rs1, uimmlog2xlen:$shamt), (XLenVT 1))>;
+def : Pat<(seteq (and GPR:$rs1, SingleBitSetMask:$mask), 0),
+          (BEXTI (XORI GPR:$rs1, -1), SingleBitSetMask:$mask)>;
 
 def : Pat<(or GPR:$r, BSETINVTwoBitsMask:$i),
           (BSETI (BSETI GPR:$r, (TrailingZerosXForm BSETINVTwoBitsMask:$i)),

llvm/test/CodeGen/RISCV/bittest.ll

@@ -156,11 +156,17 @@ define i64 @bittest_31_i64(i64 %a) nounwind {
 ; RV32-NEXT:    li a1, 0
 ; RV32-NEXT:    ret
 ;
-; RV64-LABEL: bittest_31_i64:
-; RV64:       # %bb.0:
-; RV64-NEXT:    not a0, a0
-; RV64-NEXT:    srliw a0, a0, 31
-; RV64-NEXT:    ret
+; RV64I-LABEL: bittest_31_i64:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    not a0, a0
+; RV64I-NEXT:    srliw a0, a0, 31
+; RV64I-NEXT:    ret
+;
+; RV64ZBS-LABEL: bittest_31_i64:
+; RV64ZBS:       # %bb.0:
+; RV64ZBS-NEXT:    not a0, a0
+; RV64ZBS-NEXT:    bexti a0, a0, 31
+; RV64ZBS-NEXT:    ret
   %shr = lshr i64 %a, 31
   %not = xor i64 %shr, -1
   %and = and i64 %not, 1