[RISCV] With Zbb, fold (sext_inreg (abs X)) -> (max X, (negw X))

[craig.topper]

Summary

With Zbb, abs is expanded to (max X, neg) by default. If X has 33 or
more sign bits, we can expand it a little early using negw instead of
neg to save a sext_inreg. If X started as a 32 bit value, type
legalization would have inserted a sext before the abs so X having
33 sign bits should always be true.

Note: I've used ISD::FREEZE here since we increase the number of uses.
Our default expansion for ABS doesn't do that, but I think that's a bug.

We can't do this with custom type legalization because ISD::FREEZE
doesn't propagate sign bits so later DAG combine won't expand be
able to see optmize it.

Alives2 https://alive2.llvm.org/ce/z/Gx3RNe

[achieveartificialintelligence]

llvm/test/CodeGen/RISCV/rv64zbb.ll

976

If line 964 (i32 signext %x) takes effect, why can't we remove line 977 directly?

[craig.topper]

llvm/test/CodeGen/RISCV/rv64zbb.ll

976

The absolute value of a 32-bit value is a 33-bit number because 0xffffffff80000000 becomes 0x0000000080000000. The sign extend at the end will change it back to 0xffffffff80000000.

[spatel]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

Can the ComputeNumSignBits be an assert rather than part of the predicate?

7640–7644

If this will never be anything but a quirk of the i64 type, replace VT with MVT::i64 to make that clearer?

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

The input is sign extended if the abs was promoted by type legalization, but I think it is possible to write (i64 (sext (i32 (trunc (i64 abs X)))) in the original IR and the input would not be sign extended.

[spatel]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

Maybe I'm not understanding the pattern - is it possible to write a negative test?
If we sext_inreg from i32, does this model the transform:
https://alive2.llvm.org/ce/z/j4RdVa ?

[craig.topper]

Expand the comment to better explain what is happening.

[spatel]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

I'm still not seeing it after reading the comment/example:
ashr X, 32 -> adds 32 signbits to at least 1 existing signbit
How can this be under 33?

https://alive2.llvm.org/ce/z/Rvk__m

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

Your shift result isn't being used your src function returned %abs not %ashr https://alive2.llvm.org/ce/z/eRHZww

This is the transform I'm trying to do here

define i64 @src(i64 %x) {
  %abs = call i64 @llvm.abs.i64(i64 %x, i1 0)
  %shl = shl i64 %abs, 32
  %ashr = ashr i64 %shl, 32
  ret i64 %ashr
}

define i64 @tgt(i64 %x) {
  %f = freeze i64 %x
  %negx = sub i64 0, %f
  %shl = shl i64 %negx, 32
  %ashr = ashr i64 %shl, 32
  %max = call i64 @llvm.smax.i64(i64 %ashr, i64 %negx)
  ret i64 %max
}

It's only valid if %x has 33 sign bits.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

Oops that's not right. Give me a few minutes

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

This is the transform I'm doing here

define i64 @src(i64 %x) {
  %abs = call i64 @llvm.abs.i64(i64 %x, i1 0)
  %shl = shl i64 %abs, 32
  %ashr = ashr i64 %shl, 32
  ret i64 %ashr
}

define i64 @tgt(i64 %x) {
  %f = freeze i64 %x
  %negx = sub i64 0, %f
  %shl = shl i64 %negx, 32
  %ashr = ashr i64 %shl, 32
  %max = call i64 @llvm.smax.i64(i64 %f, i64 %ashr)
  ret i64 %max
}

But it doesn't work if %x doesn't have 33 sign bits.

[spatel]

LGTM

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

7637

Thanks - now I see it. :)

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
   setTargetDAGCombine(ISD::OR);
   setTargetDAGCombine(ISD::XOR);
   if (Subtarget.hasStdExtZbp()) {
     setTargetDAGCombine(ISD::ROTL);
     setTargetDAGCombine(ISD::ROTR);
   }
   setTargetDAGCombine(ISD::ANY_EXTEND);
   setTargetDAGCombine(ISD::INTRINSIC_WO_CHAIN);
-  if (Subtarget.hasStdExtZfh())
+  if (Subtarget.hasStdExtZfh() || Subtarget.hasStdExtZbb())
     setTargetDAGCombine(ISD::SIGN_EXTEND_INREG);
   if (Subtarget.hasStdExtF()) {
     setTargetDAGCombine(ISD::ZERO_EXTEND);
     setTargetDAGCombine(ISD::FP_TO_SINT);
     setTargetDAGCombine(ISD::FP_TO_UINT);
     setTargetDAGCombine(ISD::FP_TO_SINT_SAT);
     setTargetDAGCombine(ISD::FP_TO_UINT_SAT);
   }

 }
 
 static SDValue performXORCombine(SDNode *N, SelectionDAG &DAG) {
   // fold (xor (select cond, 0, y), x) ->
   //      (select cond, x, (xor x, y))
   return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false);
 }
 
-static SDValue performSIGN_EXTEND_INREG(SDNode *N, SelectionDAG &DAG) {
+static SDValue
+performSIGN_EXTEND_INREGCombine(SDNode *N, SelectionDAG &DAG,
+                                const RISCVSubtarget &Subtarget) {
   SDValue Src = N->getOperand(0);
+  EVT VT = N->getValueType(0);
 
   // Fold (sext_inreg (fmv_x_anyexth X), i16) -> (fmv_x_signexth X)
   if (Src.getOpcode() == RISCVISD::FMV_X_ANYEXTH &&
       cast<VTSDNode>(N->getOperand(1))->getVT().bitsGE(MVT::i16))
-    return DAG.getNode(RISCVISD::FMV_X_SIGNEXTH, SDLoc(N), N->getValueType(0),
+    return DAG.getNode(RISCVISD::FMV_X_SIGNEXTH, SDLoc(N), VT,
                        Src.getOperand(0));
 
+  // Fold (i64 (sext_inreg (abs X), i32)) ->
+  // (i64 (smax (sext_inreg (neg X), i32), X)) if X has more than 32 sign bits.
+  // The (sext_inreg (neg X), i32) will be selected to negw by isel. This
+  // pattern occurs after type legalization of (i32 (abs X)) on RV64 if the user
+  // of the (i32 (abs X)) is a sext or setcc or something else that causes type
+  // legalization to add a sext_inreg after the abs. The (i32 (abs X)) will have
+  // been type legalized to (i64 (abs (sext_inreg X, i32))), but the sext_inreg
+  // may get combined into an earlier operation so we need to use
+  // ComputeNumSignBits.
+  // NOTE: (i64 (sext_inreg (abs X), i32)) can also be created for
+  // (i64 (ashr (shl (abs X), 32), 32)) without any type legalization so
+  // we can't assume that X has 33 sign bits. We must check.
+  if (Subtarget.hasStdExtZbb() && Subtarget.is64Bit() &&
+      Src.getOpcode() == ISD::ABS && Src.hasOneUse() && VT == MVT::i64 &&
+      cast<VTSDNode>(N->getOperand(1))->getVT() == MVT::i32 &&
+      DAG.ComputeNumSignBits(Src.getOperand(0)) > 32) {
+    SDLoc DL(N);
+    SDValue Freeze = DAG.getFreeze(Src.getOperand(0));
+    SDValue Neg =
+        DAG.getNode(ISD::SUB, DL, VT, DAG.getConstant(0, DL, MVT::i64), Freeze);
+    Neg = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, Neg,
+                      DAG.getValueType(MVT::i32));
+    return DAG.getNode(ISD::SMAX, DL, MVT::i64, Freeze, Neg);
+  }
+
   return SDValue();
 }
 
 // Attempt to turn ANY_EXTEND into SIGN_EXTEND if the input to the ANY_EXTEND
 // has users that require SIGN_EXTEND and the SIGN_EXTEND can be done for free
 // by an instruction like ADDW/SUBW/MULW. Without this the ANY_EXTEND would be
 // removed during type legalization leaving an ADD/SUB/MUL use that won't use
 // ADDW/SUBW/MULW.
case ISD::SUB:
     return performSUBCombine(N, DAG);
   case ISD::AND:
     return performANDCombine(N, DAG);
   case ISD::OR:
     return performORCombine(N, DAG, Subtarget);
   case ISD::XOR:
     return performXORCombine(N, DAG);
   case ISD::SIGN_EXTEND_INREG:
-    return performSIGN_EXTEND_INREG(N, DAG);
+    return performSIGN_EXTEND_INREGCombine(N, DAG, Subtarget);
   case ISD::ANY_EXTEND:
     return performANY_EXTENDCombine(N, DCI, Subtarget);
   case ISD::ZERO_EXTEND:
     // Fold (zero_extend (fp_to_uint X)) to prevent forming fcvt+zexti32 during
     // type legalization. This is safe because fp_to_uint produces poison if
     // it overflows.
     if (N->getValueType(0) == MVT::i64 && Subtarget.is64Bit()) {
       SDValue Src = N->getOperand(0);

llvm/test/CodeGen/RISCV/rv64zbb.ll

 ; 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:    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:    negw a1, a0
 ; RV64ZBB-NEXT:    max a0, a0, a1
-; RV64ZBB-NEXT:    sext.w a0, a0
 ; RV64ZBB-NEXT:    ret
   %abs = tail call i32 @llvm.abs.i32(i32 %x, i1 true)
   ret i32 %abs
 }
 
 declare i64 @llvm.abs.i64(i64, i1 immarg)
 
 define i64 @abs_i64(i64 %x) {