[RISCV] Custom type legalize i32 ISD::ABS on RV64 without Zbb.

[craig.topper]

Summary

Default type legalization will create sext_inreg+abs, but we may
not be able to remove the sext_inreg.

Instead this patch expands abs during type legalization to
Y = sraiw X, 31; subw(xor X, Y), Y) which doesn't require the input
to be sign extended.

This gives a big improvement for some neg-abs tests where the
abs is used more than the the neg. Previously the abs was expanded
a different way before and after type legalization. Now they are
expanded in a similar way enabling more CSE.

[asb]

LGTM, thanks.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (Subtarget.is64Bit()) {
       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.is64Bit())
+      setOperationAction(ISD::ABS, MVT::i32, Custom);
   }
 
   if (Subtarget.hasStdExtZbt()) {
     setOperationAction(ISD::FSHL, XLenVT, Custom);
     setOperationAction(ISD::FSHR, XLenVT, Custom);
     setOperationAction(ISD::SELECT, XLenVT, Legal);
 
     if (Subtarget.is64Bit()) {
if (Subtarget.hasStdExtZbb()) {
       return;
     }
 
     // Without Zbb, expand to UADDO/USUBO+select which will trigger our custom
     // promotion for UADDO/USUBO.
     Results.push_back(expandAddSubSat(N, DAG));
     return;
   }
+  case ISD::ABS: {
+    assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
+           "Unexpected custom legalisation");
+          DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::i64, N->getOperand(0));
+
+    // Expand abs to Y = (sraiw X, 31); subw(xor(X, Y), Y)
+
+    SDValue Src = DAG.getNode(ISD::ANY_EXTEND, DL, MVT::i64, N->getOperand(0));
+
+    // Freeze the source so we can increase it's use count.
+    Src = DAG.getFreeze(Src);
+
+    // Copy sign bit to all bits using the sraiw pattern.
+    SDValue SignFill = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, Src,
+                                   DAG.getValueType(MVT::i32));
+    SignFill = DAG.getNode(ISD::SRA, DL, MVT::i64, SignFill,
+                           DAG.getConstant(31, DL, MVT::i64));
+
+    SDValue NewRes = DAG.getNode(ISD::XOR, DL, MVT::i64, Src, SignFill);
+    NewRes = DAG.getNode(ISD::SUB, DL, MVT::i64, NewRes, SignFill);
+
+    // NOTE: The result is only required to be anyextended, but sext is
+    // consistent with type legalization of sub.
+    NewRes = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, NewRes,
+                         DAG.getValueType(MVT::i32));
+    Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, NewRes));
+    return;
+  }
   case ISD::BITCAST: {
     EVT VT = N->getValueType(0);
     assert(VT.isInteger() && !VT.isVector() && "Unexpected VT!");
     SDValue Op0 = N->getOperand(0);
     EVT Op0VT = Op0.getValueType();
     MVT XLenVT = Subtarget.getXLenVT();
     if (VT == MVT::i16 && Op0VT == MVT::f16 && Subtarget.hasStdExtZfh()) {
       SDValue FPConv = DAG.getNode(RISCVISD::FMV_X_ANYEXTH, DL, XLenVT, Op0);

llvm/test/CodeGen/RISCV/neg-abs.ll

 ; RV32IBT-NEXT:    xor a3, a0, a2
 ; RV32IBT-NEXT:    sub a0, a2, a3
 ; RV32IBT-NEXT:    sub a2, a3, a2
 ; RV32IBT-NEXT:    sw a2, 0(a1)
 ; RV32IBT-NEXT:    ret
 ;
 ; RV64I-LABEL: neg_abs32_multiuse:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sext.w a3, a0
 ; RV64I-NEXT:    sraiw a2, a0, 31
-; RV64I-NEXT:    xor a4, a0, a2
-; RV64I-NEXT:    subw a2, a2, a4
-; RV64I-NEXT:    srai a3, a3, 63
-; RV64I-NEXT:    xor a0, a0, a3
-; RV64I-NEXT:    subw a0, a0, a3
-; RV64I-NEXT:    sw a0, 0(a1)
-; RV64I-NEXT:    mv a0, a2
+; RV64I-NEXT:    xor a3, a0, a2
+; RV64I-NEXT:    subw a0, a2, a3
+; RV64I-NEXT:    subw a2, a3, a2
+; RV64I-NEXT:    sw a2, 0(a1)
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: neg_abs32_multiuse:
 ; RV64ZBB:       # %bb.0:
 ; RV64ZBB-NEXT:    sext.w a2, a0
 ; RV64ZBB-NEXT:    sraiw a3, a0, 31
 ; RV64ZBB-NEXT:    xor a0, a0, a3
 ; RV64ZBB-NEXT:    subw a0, a3, a0
 ; RV64ZBB-NEXT:    neg a3, a2
 ; RV64ZBB-NEXT:    max a2, a2, a3
 ; RV64ZBB-NEXT:    sw a2, 0(a1)
 ; RV64ZBB-NEXT:    ret
 ;
 ; RV64IBT-LABEL: neg_abs32_multiuse:
 ; RV64IBT:       # %bb.0:
-; RV64IBT-NEXT:    sext.w a3, a0
 ; RV64IBT-NEXT:    sraiw a2, a0, 31
-; RV64IBT-NEXT:    xor a4, a0, a2
-; RV64IBT-NEXT:    subw a2, a2, a4
-; RV64IBT-NEXT:    srai a3, a3, 63
-; RV64IBT-NEXT:    xor a0, a0, a3
-; RV64IBT-NEXT:    subw a0, a0, a3
-; RV64IBT-NEXT:    sw a0, 0(a1)
-; RV64IBT-NEXT:    mv a0, a2
+; RV64IBT-NEXT:    xor a3, a0, a2
+; RV64IBT-NEXT:    subw a0, a2, a3
+; RV64IBT-NEXT:    subw a2, a3, a2
+; RV64IBT-NEXT:    sw a2, 0(a1)
 ; RV64IBT-NEXT:    ret
   %abs = tail call i32 @llvm.abs.i32(i32 %x, i1 true)
   store i32 %abs, i32* %y
   %neg = sub nsw i32 0, %abs
   ret i32 %neg
 }
 
 define i64 @neg_abs64_multiuse(i64 %x, i64* %y) {

llvm/test/CodeGen/RISCV/rv64zbb.ll

; RV64ZBB-NEXT:    ret
   ret i64 %cond
 }
 
 declare i32 @llvm.abs.i32(i32, i1 immarg)
 
 define i32 @abs_i32(i32 %x) {
 ; RV64I-LABEL: abs_i32:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sext.w a0, a0
-; RV64I-NEXT:    srai a1, a0, 63
+; RV64I-NEXT:    sraiw a1, a0, 31
 ; RV64I-NEXT:    xor a0, a0, a1
-; RV64I-NEXT:    sub a0, a0, a1
+; RV64I-NEXT:    subw a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: abs_i32:
 ; RV64ZBB:       # %bb.0:
 ; RV64ZBB-NEXT:    sext.w a0, a0
 ; RV64ZBB-NEXT:    neg a1, a0
 ; RV64ZBB-NEXT:    max a0, a0, a1
 ; RV64ZBB-NEXT:    ret
   %abs = tail call i32 @llvm.abs.i32(i32 %x, i1 true)
   ret i32 %abs
 }
 
 ; FIXME: We can remove the sext.w on RV64ZBB by using negw.
 define signext i32 @abs_i32_sext(i32 signext %x) {
 ; RV64I-LABEL: abs_i32_sext:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    srai a1, a0, 63
+; RV64I-NEXT:    sraiw a1, a0, 31
 ; RV64I-NEXT:    xor a0, a0, a1
 ; RV64I-NEXT:    subw a0, a0, a1
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBB-LABEL: abs_i32_sext:
 ; RV64ZBB:       # %bb.0:
 ; RV64ZBB-NEXT:    neg a1, a0
 ; RV64ZBB-NEXT:    max a0, a0, a1