[RISCV] Add RISCVISD opcodes for CLZW and CTZW.

[craig.topper]

Summary

Our CLZW isel pattern is quite easily broken by surrounding code
preventing it from matching sometimes. This usually results in
failing to remove the and X, 0xffffffff inserted by type
legalization. The add with -32 that type legalization also inserts
often get combined into other add/sub nodes so that doesn't
usually result in extra code when we don't use clzw.

CTTZ is seems to less fragile, but I wanted to keep it consistent
with CTLZ.

[HsiangKai]

LGTM.

[asb]

LGTM, this seems the right direction to go in.

In theory I much prefer patterns in .td files wherever possible (shorter pattern description, harder for silly logic mistakes to hide within boilerplate etc). But the reality is obviously they can become brittle if above a certain level of complexity, as with many RV64 '*W' patterns.

llvm/lib/Target/RISCV/RISCVISelLowering.h

enum NodeType : unsigned {
   // by 0 or overflow (divw) like their target independent counterparts.
   DIVW,
   DIVUW,
   REMUW,
   // RV64IB rotates, directly matching the semantics of the named RISC-V
   // instructions.
   ROLW,
   RORW,
+  // RV64IZbb bit counting instructions directly matching the semantics of the
+  // named RISC-V instructions.
+  CLZW,
+  CTZW,
   // RV64IB/RV32IB funnel shifts, with the semantics of the named RISC-V
   // instructions, but the same operand order as fshl/fshr intrinsics.
   FSR,
   FSL,
   // RV64IB funnel shifts, with the semantics of the named RISC-V instructions,
   // but the same operand order as fshl/fshr intrinsics.
   FSRW,
   FSLW,

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

setOperationAction(ISD::BSWAP, XLenVT,
                        Subtarget.hasStdExtZbb() ? Legal : Expand);
   }
 
   if (Subtarget.hasStdExtZbb()) {
     setOperationAction(ISD::SMIN, XLenVT, Legal);
     setOperationAction(ISD::SMAX, XLenVT, Legal);
     setOperationAction(ISD::UMIN, XLenVT, Legal);
     setOperationAction(ISD::UMAX, XLenVT, Legal);
+
+    if (Subtarget.is64Bit()) {
+      setOperationAction(ISD::CTTZ, MVT::i32, Custom);
+      setOperationAction(ISD::CTTZ_ZERO_UNDEF, MVT::i32, Custom);
+      setOperationAction(ISD::CTLZ, MVT::i32, Custom);
+      setOperationAction(ISD::CTLZ_ZERO_UNDEF, MVT::i32, Custom);
+    }
   } else {
     setOperationAction(ISD::CTTZ, XLenVT, Expand);
     setOperationAction(ISD::CTLZ, XLenVT, Expand);
     setOperationAction(ISD::CTPOP, XLenVT, Expand);
   }
 
   if (Subtarget.hasStdExtZbt()) {
     setOperationAction(ISD::FSHL, XLenVT, Custom);
case ISD::SRL:
     Results.push_back(customLegalizeToWOp(N, DAG));
     break;
   case ISD::ROTL:
   case ISD::ROTR:
     assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
            "Unexpected custom legalisation");
     Results.push_back(customLegalizeToWOp(N, DAG));
     break;
+  case ISD::CTTZ:
+  case ISD::CTTZ_ZERO_UNDEF:
+  case ISD::CTLZ:
+  case ISD::CTLZ_ZERO_UNDEF: {
+    assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
+           "Unexpected custom legalisation");
+
+    SDValue NewOp0 =
+        DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i64, N->getOperand(0));
+    bool IsCTZ =
+        N->getOpcode() == ISD::CTTZ || N->getOpcode() == ISD::CTTZ_ZERO_UNDEF;
+    unsigned Opc = IsCTZ ? RISCVISD::CTZW : RISCVISD::CLZW;
+    SDValue Res = DAG.getNode(Opc, DL, MVT::i64, NewOp0);
+    Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, Res));
+    return;
+  }
   case ISD::SDIV:
   case ISD::UDIV:
   case ISD::UREM: {
     MVT VT = N->getSimpleValueType(0);
     assert((VT == MVT::i8 || VT == MVT::i16 || VT == MVT::i32) &&
            Subtarget.is64Bit() && Subtarget.hasStdExtM() &&
            "Unexpected custom legalisation");
     if (N->getOperand(0).getOpcode() == ISD::Constant ||
case RISCVISD::RORW: {
     if (SimplifyDemandedBits(N->getOperand(0), LHSMask, DCI) ||
         SimplifyDemandedBits(N->getOperand(1), RHSMask, DCI)) {
       if (N->getOpcode() != ISD::DELETED_NODE)
         DCI.AddToWorklist(N);
       return SDValue(N, 0);
     }
     break;
   }
+  case RISCVISD::CLZW:
+  case RISCVISD::CTZW: {
+    // Only the lower 32 bits of the first operand are read
+    SDValue Op0 = N->getOperand(0);
+    APInt Mask = APInt::getLowBitsSet(Op0.getValueSizeInBits(), 32);
+    if (SimplifyDemandedBits(Op0, Mask, DCI)) {
+      if (N->getOpcode() != ISD::DELETED_NODE)
+        DCI.AddToWorklist(N);
+      return SDValue(N, 0);
+    }
+    break;
+  }
   case RISCVISD::FSL:
   case RISCVISD::FSR: {
     // Only the lower log2(Bitwidth)+1 bits of the the shift amount are read.
     SDValue ShAmt = N->getOperand(2);
     unsigned BitWidth = ShAmt.getValueSizeInBits();
     assert(isPowerOf2_32(BitWidth) && "Unexpected bit width");
     APInt ShAmtMask(BitWidth, (BitWidth * 2) - 1);
     if (SimplifyDemandedBits(ShAmt, ShAmtMask, DCI)) {
case RISCVISD::DIVUW: {
     Known = DAG.computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
     Known2 = DAG.computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
     // We only care about the lower 32 bits.
     Known = KnownBits::udiv(Known.trunc(32), Known2.trunc(32));
     // Restore the original width by sign extending.
     Known = Known.sext(BitWidth);
     break;
   }
+  case RISCVISD::CTZW: {
+    KnownBits Known2 = DAG.computeKnownBits(Op.getOperand(0), Depth + 1);
+    unsigned PossibleTZ = Known2.trunc(32).countMaxTrailingZeros();
+    unsigned LowBits = Log2_32(PossibleTZ) + 1;
+    Known.Zero.setBitsFrom(LowBits);
+    break;
+  }
+  case RISCVISD::CLZW: {
+    KnownBits Known2 = DAG.computeKnownBits(Op.getOperand(0), Depth + 1);
+    unsigned PossibleLZ = Known2.trunc(32).countMaxLeadingZeros();
+    unsigned LowBits = Log2_32(PossibleLZ) + 1;
+    Known.Zero.setBitsFrom(LowBits);
+    break;
+  }
   case RISCVISD::READ_VLENB:
     // We assume VLENB is at least 8 bytes.
     // FIXME: The 1.0 draft spec defines minimum VLEN as 128 bits.
     Known.Zero.setLowBits(3);
     break;
   }
 }
 
#define NODE_NAME_CASE(NODE)                                                   \
   NODE_NAME_CASE(SLLW)
   NODE_NAME_CASE(SRAW)
   NODE_NAME_CASE(SRLW)
   NODE_NAME_CASE(DIVW)
   NODE_NAME_CASE(DIVUW)
   NODE_NAME_CASE(REMUW)
   NODE_NAME_CASE(ROLW)
   NODE_NAME_CASE(RORW)
+  NODE_NAME_CASE(CLZW)
+  NODE_NAME_CASE(CTZW)
   NODE_NAME_CASE(FSLW)
   NODE_NAME_CASE(FSRW)
   NODE_NAME_CASE(FSL)
   NODE_NAME_CASE(FSR)
   NODE_NAME_CASE(FMV_H_X)
   NODE_NAME_CASE(FMV_X_ANYEXTH)
   NODE_NAME_CASE(FMV_W_X_RV64)
   NODE_NAME_CASE(FMV_X_ANYEXTW_RV64)

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

 // ratified yet.
 //
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operand and SDNode transformation definitions.
 //===----------------------------------------------------------------------===//
 
+def riscv_clzw : SDNode<"RISCVISD::CLZW", SDTIntUnaryOp>;
+def riscv_ctzw : SDNode<"RISCVISD::CTZW", SDTIntUnaryOp>;
 def riscv_rolw : SDNode<"RISCVISD::ROLW", SDTIntShiftOp>;
 def riscv_rorw : SDNode<"RISCVISD::RORW", SDTIntShiftOp>;
 def riscv_fslw : SDNode<"RISCVISD::FSLW", SDTIntShiftDOp>;
 def riscv_fsrw : SDNode<"RISCVISD::FSRW", SDTIntShiftDOp>;
 def riscv_fsl  : SDNode<"RISCVISD::FSL",  SDTIntShiftDOp>;
 def riscv_fsr  : SDNode<"RISCVISD::FSR",  SDTIntShiftDOp>;
 
 def UImmLog2XLenHalfAsmOperand : AsmOperandClass {
def : Pat<(i64 (riscv_fsrw GPR:$rs3, GPR:$rs1, GPR:$rs2)),
           (FSRW GPR:$rs1, GPR:$rs2, GPR:$rs3)>;
 def : Pat<(i64 (riscv_fsrw GPR:$rs3, GPR:$rs1, uimm5:$shamt)),
           (FSRIW GPR:$rs1, GPR:$rs3, uimm5:$shamt)>;
 def : Pat<(i64 (riscv_fslw GPR:$rs3, GPR:$rs1, uimm5:$shamt)),
           (FSRIW GPR:$rs1, GPR:$rs3, (ImmSubFrom32 uimm5:$shamt))>;
 } // Predicates = [HasStdExtZbt, IsRV64]
 
 let Predicates = [HasStdExtZbb, IsRV64] in {
-def : Pat<(i64 (add (ctlz (and GPR:$rs1, 0xFFFFFFFF)), -32)),
-          (CLZW GPR:$rs1)>;
-// computeKnownBits can't figure out that the and mask on the add result is
-// unnecessary so we need to pattern match it away.
-def : Pat<(i64 (and (add (ctlz (and GPR:$rs1, 0xFFFFFFFF)), -32), 0xFFFFFFFF)),
-          (CLZW GPR:$rs1)>;
-def : Pat<(i64 (cttz (or GPR:$rs1, 0x100000000))),
-          (CTZW GPR:$rs1)>;
+def : Pat<(i64 (riscv_clzw GPR:$rs1)), (CLZW GPR:$rs1)>;
+def : Pat<(i64 (riscv_ctzw GPR:$rs1)), (CTZW GPR:$rs1)>;
 def : Pat<(i64 (ctpop (and GPR:$rs1, 0xFFFFFFFF))), (CPOPW GPR:$rs1)>;
 } // Predicates = [HasStdExtZbb, IsRV64]
 
 let Predicates = [HasStdExtZbp, IsRV64] in {
 def : Pat<(i64 (sext_inreg (or (shl GPR:$rs2, (i64 16)),
                                (and GPR:$rs1, 0x000000000000FFFF)),
                            i32)),
           (PACKW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(i64 (or (and (assertsexti32 GPR:$rs2), 0xFFFFFFFFFFFF0000),
                    (srl (and GPR:$rs1, 0xFFFFFFFF), (i64 16)))),
           (PACKUW GPR:$rs1, GPR:$rs2)>;
 } // Predicates = [HasStdExtZbp, IsRV64]

llvm/test/CodeGen/RISCV/rv64zbb.ll

 ; RV64I-NEXT:    addi a1, zero, 31
 ; RV64I-NEXT:    sub a0, a1, a0
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: log2_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    zext.w a0, a0
-; RV64IB-NEXT:    clz a0, a0
-; RV64IB-NEXT:    addi a1, zero, 63
+; RV64IB-NEXT:    clzw a0, a0
+; RV64IB-NEXT:    addi a1, zero, 31
 ; RV64IB-NEXT:    sub a0, a1, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: log2_i32:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    slli a0, a0, 32
-; RV64IBB-NEXT:    srli a0, a0, 32
-; RV64IBB-NEXT:    clz a0, a0
-; RV64IBB-NEXT:    addi a1, zero, 63
+; RV64IBB-NEXT:    clzw a0, a0
+; RV64IBB-NEXT:    addi a1, zero, 31
 ; RV64IBB-NEXT:    sub a0, a1, a0
 ; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.ctlz.i32(i32 %a, i1 false)
   %2 = sub i32 31, %1
   ret i32 %2
 }
 
 define signext i32 @log2_ceil_i32(i32 signext %a) nounwind {

 ; RV64I-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: log2_ceil_i32:
 ; RV64IB:       # %bb.0:
 ; RV64IB-NEXT:    addi a0, a0, -1
-; RV64IB-NEXT:    zext.w a0, a0
-; RV64IB-NEXT:    clz a0, a0
-; RV64IB-NEXT:    addi a1, zero, 64
+; RV64IB-NEXT:    clzw a0, a0
+; RV64IB-NEXT:    addi a1, zero, 32
 ; RV64IB-NEXT:    sub a0, a1, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: log2_ceil_i32:
 ; RV64IBB:       # %bb.0:
 ; RV64IBB-NEXT:    addi a0, a0, -1
-; RV64IBB-NEXT:    slli a0, a0, 32
-; RV64IBB-NEXT:    srli a0, a0, 32
-; RV64IBB-NEXT:    clz a0, a0
-; RV64IBB-NEXT:    addi a1, zero, 64
+; RV64IBB-NEXT:    clzw a0, a0
+; RV64IBB-NEXT:    addi a1, zero, 32
 ; RV64IBB-NEXT:    sub a0, a1, a0
 ; RV64IBB-NEXT:    ret
   %1 = sub i32 %a, 1
   %2 = call i32 @llvm.ctlz.i32(i32 %1, i1 false)
   %3 = sub i32 32, %2
   ret i32 %3
 }
 

 ; RV64I-NEXT:  .LBB4_3: # %cond.end
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: ctlz_lshr_i32:
 ; RV64IB:       # %bb.0:
 ; RV64IB-NEXT:    srliw a0, a0, 1
-; RV64IB-NEXT:    clz a0, a0
-; RV64IB-NEXT:    addi a0, a0, -32
+; RV64IB-NEXT:    clzw a0, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: ctlz_lshr_i32:
 ; RV64IBB:       # %bb.0:
 ; RV64IBB-NEXT:    srliw a0, a0, 1
-; RV64IBB-NEXT:    clz a0, a0
-; RV64IBB-NEXT:    addi a0, a0, -32
+; RV64IBB-NEXT:    clzw a0, a0
 ; RV64IBB-NEXT:    ret
   %1 = lshr i32 %a, 1
   %2 = call i32 @llvm.ctlz.i32(i32 %1, i1 false)
   ret i32 %2
 }
 
 declare i64 @llvm.ctlz.i64(i64, i1)
 

 ; RV64I-NEXT:    call __muldi3@plt
 ; RV64I-NEXT:    srli a0, a0, 56
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: cttz_zero_undef_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    ctz a0, a0
+; RV64IB-NEXT:    ctzw a0, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: cttz_zero_undef_i32:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    ctz a0, a0
+; RV64IBB-NEXT:    ctzw a0, a0
 ; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.cttz.i32(i32 %a, i1 true)
   ret i32 %1
 }
 
 define signext i32 @findFirstSet_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: findFirstSet_i32:
 ; RV64I:       # %bb.0:

 ; RV64I-NEXT:  .LBB8_2:
 ; RV64I-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: findFirstSet_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    ctz a1, a0
+; RV64IB-NEXT:    ctzw a1, a0
 ; RV64IB-NEXT:    addi a2, zero, -1
 ; RV64IB-NEXT:    cmov a0, a0, a1, a2
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: findFirstSet_i32:
 ; RV64IBB:       # %bb.0:
 ; RV64IBB-NEXT:    mv a1, a0
 ; RV64IBB-NEXT:    addi a0, zero, -1
 ; RV64IBB-NEXT:    beqz a1, .LBB8_2
 ; RV64IBB-NEXT:  # %bb.1:
-; RV64IBB-NEXT:    ctz a0, a1
+; RV64IBB-NEXT:    ctzw a0, a1
 ; RV64IBB-NEXT:  .LBB8_2:
 ; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.cttz.i32(i32 %a, i1 true)
   %2 = icmp eq i32 %a, 0
   %3 = select i1 %2, i32 -1, i32 %1
   ret i32 %3
 }
 

 ; RV64I-NEXT:  .LBB9_2:
 ; RV64I-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: ffs_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    ctz a1, a0
+; RV64IB-NEXT:    ctzw a1, a0
 ; RV64IB-NEXT:    addi a1, a1, 1
 ; RV64IB-NEXT:    cmov a0, a0, a1, zero
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: ffs_i32:
 ; RV64IBB:       # %bb.0:
 ; RV64IBB-NEXT:    mv a1, a0
 ; RV64IBB-NEXT:    mv a0, zero
 ; RV64IBB-NEXT:    beqz a1, .LBB9_2
 ; RV64IBB-NEXT:  # %bb.1:
-; RV64IBB-NEXT:    ctz a0, a1
+; RV64IBB-NEXT:    ctzw a0, a1
 ; RV64IBB-NEXT:    addi a0, a0, 1
 ; RV64IBB-NEXT:  .LBB9_2:
 ; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.cttz.i32(i32 %a, i1 true)
   %2 = add i32 %1, 1
   %3 = icmp eq i32 %a, 0
   %4 = select i1 %3, i32 0, i32 %2
   ret i32 %4