[RISCV] Custom legalize i32 saddo/ssubo on RV64 to return a sign extended value for the data result. (#172112)

This is consistent with how we handle regular ADD/SUB and helps with
computeNumSignBits optimizations.

Fixes #172089

Author: topperc

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -358,7 +358,7 @@ RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
       setOperationAction(
           {ISD::SADDSAT, ISD::SSUBSAT, ISD::UADDSAT, ISD::USUBSAT}, MVT::i32,
           Custom);
-    setOperationAction(ISD::SADDO, MVT::i32, Custom);
+    setOperationAction({ISD::SADDO, ISD::SSUBO}, MVT::i32, Custom);
   }
   if (!Subtarget.hasStdExtZmmul()) {
     setOperationAction({ISD::MUL, ISD::MULHS, ISD::MULHU}, XLenVT, Expand);
@@ -15060,35 +15060,42 @@ void RISCVTargetLowering::ReplaceNodeResults(SDNode *N,
     Results.push_back(customLegalizeToWOp(N, DAG, ExtOpc));
     break;
   }
-  case ISD::SADDO: {
+  case ISD::SADDO:
+  case ISD::SSUBO: {
     assert(N->getValueType(0) == MVT::i32 && Subtarget.is64Bit() &&
            "Unexpected custom legalisation");
 
-    // If the RHS is a constant, we can simplify ConditionRHS below. Otherwise
-    // use the default legalization.
-    if (!isa<ConstantSDNode>(N->getOperand(1)))
-      return;
-
+    // This is similar to the default legalization, but we return the
+    // sext_inreg instead of the add/sub.
+    bool IsAdd = N->getOpcode() == ISD::SADDO;
     SDValue LHS = DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::i64, N->getOperand(0));
     SDValue RHS = DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::i64, N->getOperand(1));
-    SDValue Res = DAG.getNode(ISD::ADD, DL, MVT::i64, LHS, RHS);
-    Res = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, Res,
-                      DAG.getValueType(MVT::i32));
+    SDValue Op =
+        DAG.getNode(IsAdd ? ISD::ADD : ISD::SUB, DL, MVT::i64, LHS, RHS);
+    SDValue Res = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i64, Op,
+                              DAG.getValueType(MVT::i32));
 
-    SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
+    SDValue Overflow;
 
-    // For an addition, the result should be less than one of the operands (LHS)
-    // if and only if the other operand (RHS) is negative, otherwise there will
-    // be overflow.
-    // For a subtraction, the result should be less than one of the operands
-    // (LHS) if and only if the other operand (RHS) is (non-zero) positive,
-    // otherwise there will be overflow.
-    EVT OType = N->getValueType(1);
-    SDValue ResultLowerThanLHS = DAG.getSetCC(DL, OType, Res, LHS, ISD::SETLT);
-    SDValue ConditionRHS = DAG.getSetCC(DL, OType, RHS, Zero, ISD::SETLT);
+    // If the RHS is a constant, we can simplify ConditionRHS below. Otherwise
+    // use the default legalization.
+    if (IsAdd && isa<ConstantSDNode>(N->getOperand(1))) {
+      SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
+
+      // For an addition, the result should be less than one of the operands
+      // (LHS) if and only if the other operand (RHS) is negative, otherwise
+      // there will be overflow.
+      EVT OType = N->getValueType(1);
+      SDValue ResultLowerThanLHS =
+          DAG.getSetCC(DL, OType, Res, LHS, ISD::SETLT);
+      SDValue ConditionRHS = DAG.getSetCC(DL, OType, RHS, Zero, ISD::SETLT);
+
+      Overflow =
+          DAG.getNode(ISD::XOR, DL, OType, ConditionRHS, ResultLowerThanLHS);
+    } else {
+      Overflow = DAG.getSetCC(DL, N->getValueType(1), Res, Op, ISD::SETNE);
+    }
 
-    SDValue Overflow =
-        DAG.getNode(ISD::XOR, DL, OType, ConditionRHS, ResultLowerThanLHS);
     Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, Res));
     Results.push_back(Overflow);
     return;

llvm/test/CodeGen/RISCV/sadd_sat.ll

@@ -21,16 +21,15 @@ define signext i32 @func(i32 signext %x, i32 signext %y) nounwind {
 ;
 ; RV64I-LABEL: func:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    add a2, a0, a1
+; RV64I-NEXT:    mv a2, a0
 ; RV64I-NEXT:    addw a0, a0, a1
-; RV64I-NEXT:    beq a0, a2, .LBB0_2
+; RV64I-NEXT:    add a1, a2, a1
+; RV64I-NEXT:    beq a0, a1, .LBB0_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    srli a0, a0, 31
-; RV64I-NEXT:    li a1, 1
-; RV64I-NEXT:    slli a1, a1, 31
-; RV64I-NEXT:    xor a2, a0, a1
+; RV64I-NEXT:    sraiw a0, a1, 31
+; RV64I-NEXT:    lui a1, 524288
+; RV64I-NEXT:    xor a0, a0, a1
 ; RV64I-NEXT:  .LBB0_2:
-; RV64I-NEXT:    sext.w a0, a2
 ; RV64I-NEXT:    ret
 ;
 ; RV64IZbb-LABEL: func:

llvm/test/CodeGen/RISCV/sadd_sat_plus.ll

@@ -22,13 +22,13 @@ define i32 @func32(i32 %x, i32 %y, i32 %z) nounwind {
 ;
 ; RV64I-LABEL: func32:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sext.w a0, a0
 ; RV64I-NEXT:    mulw a1, a1, a2
-; RV64I-NEXT:    addw a2, a0, a1
-; RV64I-NEXT:    add a0, a0, a1
-; RV64I-NEXT:    beq a2, a0, .LBB0_2
+; RV64I-NEXT:    sext.w a2, a0
+; RV64I-NEXT:    addw a0, a2, a1
+; RV64I-NEXT:    add a1, a2, a1
+; RV64I-NEXT:    beq a0, a1, .LBB0_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    sraiw a0, a0, 31
+; RV64I-NEXT:    sraiw a0, a1, 31
 ; RV64I-NEXT:    lui a1, 524288
 ; RV64I-NEXT:    xor a0, a0, a1
 ; RV64I-NEXT:  .LBB0_2:

llvm/test/CodeGen/RISCV/ssub_sat.ll

@@ -21,16 +21,15 @@ define signext i32 @func(i32 signext %x, i32 signext %y) nounwind {
 ;
 ; RV64I-LABEL: func:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sub a2, a0, a1
+; RV64I-NEXT:    mv a2, a0
 ; RV64I-NEXT:    subw a0, a0, a1
+; RV64I-NEXT:    sub a2, a2, a1
 ; RV64I-NEXT:    beq a0, a2, .LBB0_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    srli a0, a0, 31
-; RV64I-NEXT:    li a1, 1
-; RV64I-NEXT:    slli a1, a1, 31
-; RV64I-NEXT:    xor a2, a0, a1
+; RV64I-NEXT:    sraiw a0, a2, 31
+; RV64I-NEXT:    lui a1, 524288
+; RV64I-NEXT:    xor a0, a0, a1
 ; RV64I-NEXT:  .LBB0_2:
-; RV64I-NEXT:    sext.w a0, a2
 ; RV64I-NEXT:    ret
 ;
 ; RV64IZbb-LABEL: func:

llvm/test/CodeGen/RISCV/ssub_sat_plus.ll

@@ -22,13 +22,13 @@ define i32 @func32(i32 %x, i32 %y, i32 %z) nounwind {
 ;
 ; RV64I-LABEL: func32:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    sext.w a0, a0
 ; RV64I-NEXT:    mulw a1, a1, a2
-; RV64I-NEXT:    subw a2, a0, a1
-; RV64I-NEXT:    sub a0, a0, a1
-; RV64I-NEXT:    beq a2, a0, .LBB0_2
+; RV64I-NEXT:    sext.w a2, a0
+; RV64I-NEXT:    subw a0, a2, a1
+; RV64I-NEXT:    sub a1, a2, a1
+; RV64I-NEXT:    beq a0, a1, .LBB0_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    sraiw a0, a0, 31
+; RV64I-NEXT:    sraiw a0, a1, 31
 ; RV64I-NEXT:    lui a1, 524288
 ; RV64I-NEXT:    xor a0, a0, a1
 ; RV64I-NEXT:  .LBB0_2: