[RISCV] Make ctlz/cttz cheap to speculatively execute so CodeGenPrepare won't insert a zero check.

[craig.topper]

Summary

Add additional isel patterns for ctzw/clzw instructions.

[lenary]

LGTM!

llvm/lib/Target/RISCV/RISCVISelLowering.h

bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
                              unsigned AS,
                              Instruction *I = nullptr) const override;
   bool isLegalICmpImmediate(int64_t Imm) const override;
   bool isLegalAddImmediate(int64_t Imm) const override;
   bool isTruncateFree(Type *SrcTy, Type *DstTy) const override;
   bool isTruncateFree(EVT SrcVT, EVT DstVT) const override;
   bool isZExtFree(SDValue Val, EVT VT2) const override;
   bool isSExtCheaperThanZExt(EVT SrcVT, EVT DstVT) const override;
+  bool isCheapToSpeculateCttz() const override;
+  bool isCheapToSpeculateCtlz() const override;
   bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;
 
   bool hasBitPreservingFPLogic(EVT VT) const override;
 
   // Provide custom lowering hooks for some operations.
   SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

bool RISCVTargetLowering::isZExtFree(SDValue Val, EVT VT2) const {
 
   return TargetLowering::isZExtFree(Val, VT2);
 }
 
 bool RISCVTargetLowering::isSExtCheaperThanZExt(EVT SrcVT, EVT DstVT) const {
   return Subtarget.is64Bit() && SrcVT == MVT::i32 && DstVT == MVT::i64;
 }
 
+bool RISCVTargetLowering::isCheapToSpeculateCttz() const {
+  return Subtarget.hasStdExtZbb();
+}
+
+bool RISCVTargetLowering::isCheapToSpeculateCtlz() const {
+  return Subtarget.hasStdExtZbb();
+}
+
 bool RISCVTargetLowering::isFPImmLegal(const APFloat &Imm, EVT VT,
                                        bool ForCodeSize) const {
   if (VT == MVT::f32 && !Subtarget.hasStdExtF())
     return false;
   if (VT == MVT::f64 && !Subtarget.hasStdExtD())
     return false;
   if (Imm.isNegZero())
     return false;

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

def : Pat<(sext_inreg (fshl GPR:$rs3, (shl GPR:$rs1, (i64 32)),
                             uimm5:$shamt),
                       i32),
           (FSRIW GPR:$rs1, GPR:$rs3, (ImmROTL2RW uimm5:$shamt))>;
 } // Predicates = [HasStdExtZbt, IsRV64]
 
 let Predicates = [HasStdExtZbb, IsRV64] in {
 def : Pat<(add (ctlz (and GPR:$rs1, (i64 0xFFFFFFFF))), (i64 -32)),
           (CLZW GPR:$rs1)>;
-// We don't pattern-match CTZW here as it has the same pattern and result as
-// RV64 CTZ
+// 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<(and (add (ctlz (and GPR:$rs1, (i64 0xFFFFFFFF))), (i64 -32)),
+               (i64 0xFFFFFFFF)),
+          (CLZW GPR:$rs1)>;
+def : Pat<(cttz (or GPR:$rs1, (i64 0x100000000))),
+          (CTZW GPR:$rs1)>;
 def : Pat<(ctpop (and GPR:$rs1, (i64 0xFFFFFFFF))), (PCNTW GPR:$rs1)>;
 } // Predicates = [HasStdExtZbb, IsRV64]
 
 let Predicates = [HasStdExtZbbOrZbp, IsRV64] in {
 def : Pat<(sext_inreg (or (shl GPR:$rs2, (i64 16)),
                           (and GPR:$rs1, 0x000000000000FFFF)),
                       i32),
           (PACKW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(or (and (assertsexti32 GPR:$rs2), 0xFFFFFFFFFFFF0000),
               (srl (and GPR:$rs1, 0x00000000FFFF0000),
                    (i64 16))),
           (PACKUW GPR:$rs1, GPR:$rs2)>;
 } // Predicates = [HasStdExtZbbOrZbp, IsRV64]

llvm/test/CodeGen/RISCV/rv32Zbb.ll

 ; RV32I-NEXT:    addi a0, zero, 32
 ; RV32I-NEXT:  .LBB8_3: # %cond.end
 ; RV32I-NEXT:    lw ra, 12(sp)
 ; RV32I-NEXT:    addi sp, sp, 16
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: ctlz_i32:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    beqz a0, .LBB8_2
-; RV32IB-NEXT:  # %bb.1: # %cond.false
 ; RV32IB-NEXT:    clz a0, a0
 ; RV32IB-NEXT:    ret
-; RV32IB-NEXT:  .LBB8_2:
-; RV32IB-NEXT:    addi a0, zero, 32
-; RV32IB-NEXT:    ret
 ;
 ; RV32IBB-LABEL: ctlz_i32:
 ; RV32IBB:       # %bb.0:
-; RV32IBB-NEXT:    beqz a0, .LBB8_2
-; RV32IBB-NEXT:  # %bb.1: # %cond.false
 ; RV32IBB-NEXT:    clz a0, a0
 ; RV32IBB-NEXT:    ret
-; RV32IBB-NEXT:  .LBB8_2:
-; RV32IBB-NEXT:    addi a0, zero, 32
-; RV32IBB-NEXT:    ret
   %1 = call i32 @llvm.ctlz.i32(i32 %a, i1 false)
   ret i32 %1
 }
 
 declare i64 @llvm.ctlz.i64(i64, i1)
 
 define i64 @ctlz_i64(i64 %a) nounwind {
 ; RV32I-LABEL: ctlz_i64:

 ; RV32I-NEXT:    addi a0, zero, 32
 ; RV32I-NEXT:  .LBB10_3: # %cond.end
 ; RV32I-NEXT:    lw ra, 12(sp)
 ; RV32I-NEXT:    addi sp, sp, 16
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: cttz_i32:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    beqz a0, .LBB10_2
-; RV32IB-NEXT:  # %bb.1: # %cond.false
 ; RV32IB-NEXT:    ctz a0, a0
 ; RV32IB-NEXT:    ret
-; RV32IB-NEXT:  .LBB10_2:
-; RV32IB-NEXT:    addi a0, zero, 32
-; RV32IB-NEXT:    ret
 ;
 ; RV32IBB-LABEL: cttz_i32:
 ; RV32IBB:       # %bb.0:
-; RV32IBB-NEXT:    beqz a0, .LBB10_2
-; RV32IBB-NEXT:  # %bb.1: # %cond.false
 ; RV32IBB-NEXT:    ctz a0, a0
 ; RV32IBB-NEXT:    ret
-; RV32IBB-NEXT:  .LBB10_2:
-; RV32IBB-NEXT:    addi a0, zero, 32
-; RV32IBB-NEXT:    ret
   %1 = call i32 @llvm.cttz.i32(i32 %a, i1 false)
   ret i32 %1
 }
 
 declare i64 @llvm.cttz.i64(i64, i1)
 
 define i64 @cttz_i64(i64 %a) nounwind {
 ; RV32I-LABEL: cttz_i64:

llvm/test/CodeGen/RISCV/rv64Zbb.ll

 ; RV64I-NEXT:    addi a0, zero, 32
 ; RV64I-NEXT:  .LBB9_3: # %cond.end
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: ctlz_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    beqz a0, .LBB9_2
-; RV64IB-NEXT:  # %bb.1: # %cond.false
 ; RV64IB-NEXT:    clzw a0, a0
 ; RV64IB-NEXT:    ret
-; RV64IB-NEXT:  .LBB9_2:
-; RV64IB-NEXT:    addi a0, zero, 32
-; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: ctlz_i32:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    beqz a0, .LBB9_2
-; RV64IBB-NEXT:  # %bb.1: # %cond.false
 ; RV64IBB-NEXT:    clzw a0, a0
 ; RV64IBB-NEXT:    ret
-; RV64IBB-NEXT:  .LBB9_2:
-; RV64IBB-NEXT:    addi a0, zero, 32
-; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.ctlz.i32(i32 %a, i1 false)
   ret i32 %1
 }
 
 declare i64 @llvm.ctlz.i64(i64, i1)
 
 define i64 @ctlz_i64(i64 %a) nounwind {
 ; RV64I-LABEL: ctlz_i64:

 ; RV64I-NEXT:    addi a0, zero, 64
 ; RV64I-NEXT:  .LBB10_3: # %cond.end
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: ctlz_i64:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    beqz a0, .LBB10_2
-; RV64IB-NEXT:  # %bb.1: # %cond.false
 ; RV64IB-NEXT:    clz a0, a0
 ; RV64IB-NEXT:    ret
-; RV64IB-NEXT:  .LBB10_2:
-; RV64IB-NEXT:    addi a0, zero, 64
-; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: ctlz_i64:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    beqz a0, .LBB10_2
-; RV64IBB-NEXT:  # %bb.1: # %cond.false
 ; RV64IBB-NEXT:    clz a0, a0
 ; RV64IBB-NEXT:    ret
-; RV64IBB-NEXT:  .LBB10_2:
-; RV64IBB-NEXT:    addi a0, zero, 64
-; RV64IBB-NEXT:    ret
   %1 = call i64 @llvm.ctlz.i64(i64 %a, i1 false)
   ret i64 %1
 }
 
 declare i32 @llvm.cttz.i32(i32, i1)
 
 define signext i32 @cttz_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: cttz_i32:

 ; RV64I-NEXT:    addi a0, zero, 32
 ; RV64I-NEXT:  .LBB11_3: # %cond.end
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: cttz_i32:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    beqz a0, .LBB11_2
-; RV64IB-NEXT:  # %bb.1: # %cond.false
-; RV64IB-NEXT:    ctz a0, a0
-; RV64IB-NEXT:    ret
-; RV64IB-NEXT:  .LBB11_2:
-; RV64IB-NEXT:    addi a0, zero, 32
+; RV64IB-NEXT:    ctzw a0, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: cttz_i32:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    beqz a0, .LBB11_2
-; RV64IBB-NEXT:  # %bb.1: # %cond.false
-; RV64IBB-NEXT:    ctz a0, a0
-; RV64IBB-NEXT:    ret
-; RV64IBB-NEXT:  .LBB11_2:
-; RV64IBB-NEXT:    addi a0, zero, 32
+; RV64IBB-NEXT:    ctzw a0, a0
 ; RV64IBB-NEXT:    ret
   %1 = call i32 @llvm.cttz.i32(i32 %a, i1 false)
   ret i32 %1
 }
 
 declare i64 @llvm.cttz.i64(i64, i1)
 
 define i64 @cttz_i64(i64 %a) nounwind {

 ; RV64I-NEXT:    addi a0, zero, 64
 ; RV64I-NEXT:  .LBB12_3: # %cond.end
 ; RV64I-NEXT:    ld ra, 8(sp)
 ; RV64I-NEXT:    addi sp, sp, 16
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: cttz_i64:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    beqz a0, .LBB12_2
-; RV64IB-NEXT:  # %bb.1: # %cond.false
 ; RV64IB-NEXT:    ctz a0, a0
 ; RV64IB-NEXT:    ret
-; RV64IB-NEXT:  .LBB12_2:
-; RV64IB-NEXT:    addi a0, zero, 64
-; RV64IB-NEXT:    ret
 ;
 ; RV64IBB-LABEL: cttz_i64:
 ; RV64IBB:       # %bb.0:
-; RV64IBB-NEXT:    beqz a0, .LBB12_2
-; RV64IBB-NEXT:  # %bb.1: # %cond.false
 ; RV64IBB-NEXT:    ctz a0, a0
 ; RV64IBB-NEXT:    ret
-; RV64IBB-NEXT:  .LBB12_2:
-; RV64IBB-NEXT:    addi a0, zero, 64
-; RV64IBB-NEXT:    ret
   %1 = call i64 @llvm.cttz.i64(i64 %a, i1 false)
   ret i64 %1
 }
 
 declare i32 @llvm.ctpop.i32(i32)
 
 define signext i32 @ctpop_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: ctpop_i32: