[RISCV] Custom type legalize i32 bswap/bitreverse to GREVIW on RV64 with Zbp extension

[craig.topper]

Summary

Currently we require a sra to pattern match these properly in isel. If the consumer didn't need the result sign extended we'll have an srl instead of sra and fail to match.

This patch switches to custom legalizing to GREVIW using portions of D91259.

[craig.topper]

Rebase

[frasercrmck]

LGTM.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (Subtarget.is64Bit()) {
       setOperationAction(ISD::ROTL, MVT::i32, Custom);
       setOperationAction(ISD::ROTR, MVT::i32, Custom);
     }
   } else {
     setOperationAction(ISD::ROTL, XLenVT, Expand);
     setOperationAction(ISD::ROTR, XLenVT, Expand);
   }
 
-  if (!Subtarget.hasStdExtZbp())
+  if (Subtarget.hasStdExtZbp()) {
+    setOperationAction(ISD::BITREVERSE, XLenVT, Legal);
+
+    if (Subtarget.is64Bit()) {
+      setOperationAction(ISD::BITREVERSE, MVT::i32, Custom);
+      setOperationAction(ISD::BSWAP, MVT::i32, Custom);
+    }
+  } else {
     setOperationAction(ISD::BSWAP, XLenVT, Expand);
+  }
 
   if (!Subtarget.hasStdExtZbb()) {
     setOperationAction(ISD::CTTZ, XLenVT, Expand);
     setOperationAction(ISD::CTLZ, XLenVT, Expand);
     setOperationAction(ISD::CTPOP, XLenVT, Expand);
   }
 
-  if (Subtarget.hasStdExtZbp())
-    setOperationAction(ISD::BITREVERSE, XLenVT, Legal);
-
   if (Subtarget.hasStdExtZbt()) {
     setOperationAction(ISD::FSHL, XLenVT, Legal);
     setOperationAction(ISD::FSHR, XLenVT, Legal);
   }
 
   ISD::CondCode FPCCToExtend[] = {
       ISD::SETOGT, ISD::SETOGE, ISD::SETONE, ISD::SETUEQ, ISD::SETUGT,
       ISD::SETUGE, ISD::SETULT, ISD::SETULE, ISD::SETUNE, ISD::SETGT,
SDValue NewOp0 =
         DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i64, N->getOperand(0));
     SDValue NewRes =
         DAG.getNode(WOpcode, DL, MVT::i64, NewOp0, N->getOperand(1));
     // ReplaceNodeResults requires we maintain the same type for the return
     // value.
     Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, NewRes));
     break;
   }
+  case ISD::BSWAP:
+  case ISD::BITREVERSE: {
+    assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
+           Subtarget.hasStdExtZbp() && "Unexpected custom legalisation");
+    SDValue NewOp0 = DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i64,
+                                 N->getOperand(0));
+    unsigned Imm = N->getOpcode() == ISD::BITREVERSE ? 31 : 24;
+    SDValue GREVIW = DAG.getNode(RISCVISD::GREVIW, DL, MVT::i64, NewOp0,
+                                 DAG.getTargetConstant(Imm, DL,
+                                                       Subtarget.getXLenVT()));
+    // ReplaceNodeResults requires we maintain the same type for the return
+    // value.
+    Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, GREVIW));
+    break;
+  }
   }
 }
 
 // A structure to hold one of the bit-manipulation patterns below. Together, a
 // SHL and non-SHL pattern may form a bit-manipulation pair on a single source:
 //   (or (and (shl x, 1), 0xAAAAAAAA),
 //       (and (srl x, 1), 0x55555555))
 struct RISCVBitmanipPat {

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

def : Pat<(SLOIWPat GPR:$rs1, uimmlog2xlen:$shamt),
           (SLOIW GPR:$rs1, uimmlog2xlen:$shamt)>;
 def : Pat<(SROIWPat GPR:$rs1, uimmlog2xlen:$shamt),
           (SROIW GPR:$rs1, uimmlog2xlen:$shamt)>;
 } // Predicates = [HasStdExtZbb, IsRV64]
 
 let Predicates = [HasStdExtZbp, IsRV64] in {
 def : Pat<(riscv_greviw GPR:$rs1, timm:$shamt), (GREVIW GPR:$rs1, timm:$shamt)>;
 def : Pat<(riscv_gorciw GPR:$rs1, timm:$shamt), (GORCIW GPR:$rs1, timm:$shamt)>;
-
-def : Pat<(sra (bswap GPR:$rs1), (i64 32)), (GREVIW GPR:$rs1, (i64 24))>;
-def : Pat<(sra (bitreverse GPR:$rs1), (i64 32)), (GREVIW GPR:$rs1, (i64 31))>;
 } // Predicates = [HasStdExtZbp, IsRV64]
 
 let Predicates = [HasStdExtZbt, IsRV64] in {
 def : Pat<(sext_inreg (fshl GPR:$rs1, (shl GPR:$rs3, (i64 32)),
                             (and GPR:$rs2, (i64 31))),
                       i32),
           (FSLW GPR:$rs1, (ANDI GPR:$rs2, 31), GPR:$rs3)>;
 def : Pat<(sext_inreg (fshr GPR:$rs3, (shl GPR:$rs1, (i64 32)),

llvm/test/CodeGen/RISCV/rv64Zbp.ll

 ; RV64I-NEXT:    or a0, a0, a3
 ; RV64I-NEXT:    or a0, a0, a2
 ; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    sw a0, 0(a1)
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: bswap_i32_nosext:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    rev8 a0, a0
-; RV64IB-NEXT:    srli a0, a0, 32
+; RV64IB-NEXT:    greviw a0, a0, 24
 ; RV64IB-NEXT:    sw a0, 0(a1)
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBP-LABEL: bswap_i32_nosext:
 ; RV64IBP:       # %bb.0:
-; RV64IBP-NEXT:    rev8 a0, a0
-; RV64IBP-NEXT:    srli a0, a0, 32
+; RV64IBP-NEXT:    greviw a0, a0, 24
 ; RV64IBP-NEXT:    sw a0, 0(a1)
 ; RV64IBP-NEXT:    ret
   %1 = tail call i32 @llvm.bswap.i32(i32 %a)
   store i32 %1, i32* %x
   ret void
 }
 
 declare i64 @llvm.bswap.i64(i64)

 ; RV64I-NEXT:    srli a0, a0, 1
 ; RV64I-NEXT:    or a0, a0, a2
 ; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    sw a0, 0(a1)
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: bitreverse_i32_nosext:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    rev a0, a0
-; RV64IB-NEXT:    srli a0, a0, 32
+; RV64IB-NEXT:    greviw a0, a0, 31
 ; RV64IB-NEXT:    sw a0, 0(a1)
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBP-LABEL: bitreverse_i32_nosext:
 ; RV64IBP:       # %bb.0:
-; RV64IBP-NEXT:    rev a0, a0
-; RV64IBP-NEXT:    srli a0, a0, 32
+; RV64IBP-NEXT:    greviw a0, a0, 31
 ; RV64IBP-NEXT:    sw a0, 0(a1)
 ; RV64IBP-NEXT:    ret
   %1 = tail call i32 @llvm.bitreverse.i32(i32 %a)
   store i32 %1, i32* %x
   ret void
 }
 
 declare i64 @llvm.bitreverse.i64(i64)