[SelectionDAG] Support scalable splats in U(ADD|SUB)SAT combines

[frasercrmck]

Summary

This patch builds on top of D106575 in which scalable-vector splats were
supported in ISD::matchBinaryPredicate. It teaches the DAGCombiner how
to perform a variety of the pre-existing saturating add/sub combines on
scalable-vector types.

[frasercrmck]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10053

I was wondering whether now's the time to add yet another function like ISD::isBuildVectorOrSplatVector. Thoughts?

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10109

Do the opcodes need to match here?

[frasercrmck]

ensure matching opcodes

[craig.topper]

LGTM

[frasercrmck]

rebase (arcanist lost track)

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (N0.getOpcode() == ISD::SETCC) {
     }
 
     // Match VSELECTs into add with unsigned saturation.
     if (hasOperation(ISD::UADDSAT, VT)) {
       // Check if one of the arms of the VSELECT is vector with all bits set.
       // If it's on the left side invert the predicate to simplify logic below.
       SDValue Other;
       ISD::CondCode SatCC = CC;
-      if (ISD::isBuildVectorAllOnes(N1.getNode())) {
+      if (ISD::isConstantSplatVectorAllOnes(N1.getNode())) {
         Other = N2;
         SatCC = ISD::getSetCCInverse(SatCC, VT.getScalarType());
-      } else if (ISD::isBuildVectorAllOnes(N2.getNode())) {
+      } else if (ISD::isConstantSplatVectorAllOnes(N2.getNode())) {
         Other = N1;
       }
 
       if (Other && Other.getOpcode() == ISD::ADD) {
         SDValue CondLHS = LHS, CondRHS = RHS;
         SDValue OpLHS = Other.getOperand(0), OpRHS = Other.getOperand(1);
 
         // Canonicalize condition operands.
         if (SatCC == ISD::SETUGE) {
           std::swap(CondLHS, CondRHS);
           SatCC = ISD::SETULE;
         }
 
         // We can test against either of the addition operands.
         // x <= x+y ? x+y : ~0 --> uaddsat x, y
         // x+y >= x ? x+y : ~0 --> uaddsat x, y
         if (SatCC == ISD::SETULE && Other == CondRHS &&
             (OpLHS == CondLHS || OpRHS == CondLHS))
           return DAG.getNode(ISD::UADDSAT, DL, VT, OpLHS, OpRHS);
 
-        if (isa<BuildVectorSDNode>(OpRHS) && isa<BuildVectorSDNode>(CondRHS) &&
+        if (OpRHS.getOpcode() == CondRHS.getOpcode() &&
+            (OpRHS.getOpcode() == ISD::BUILD_VECTOR ||
+             OpRHS.getOpcode() == ISD::SPLAT_VECTOR) &&
             CondLHS == OpLHS) {
           // If the RHS is a constant we have to reverse the const
           // canonicalization.
           // x >= ~C ? x+C : ~0 --> uaddsat x, C
           auto MatchUADDSAT = [](ConstantSDNode *Op, ConstantSDNode *Cond) {
             return Cond->getAPIntValue() == ~Op->getAPIntValue();
           };
           if (SatCC == ISD::SETULE &&
               ISD::matchBinaryPredicate(OpRHS, CondRHS, MatchUADDSAT))
             return DAG.getNode(ISD::UADDSAT, DL, VT, OpLHS, OpRHS);
         }
       }
     }
 
     // Match VSELECTs into sub with unsigned saturation.
     if (hasOperation(ISD::USUBSAT, VT)) {
       // Check if one of the arms of the VSELECT is a zero vector. If it's on
       // the left side invert the predicate to simplify logic below.
       SDValue Other;
       ISD::CondCode SatCC = CC;
-      if (ISD::isBuildVectorAllZeros(N1.getNode())) {
+      if (ISD::isConstantSplatVectorAllZeros(N1.getNode())) {
         Other = N2;
         SatCC = ISD::getSetCCInverse(SatCC, VT.getScalarType());
-      } else if (ISD::isBuildVectorAllZeros(N2.getNode())) {
+      } else if (ISD::isConstantSplatVectorAllZeros(N2.getNode())) {
         Other = N1;
       }
 
       if (Other && Other.getNumOperands() == 2) {
         SDValue CondRHS = RHS;
         SDValue OpLHS = Other.getOperand(0), OpRHS = Other.getOperand(1);
 
         if (Other.getOpcode() == ISD::SUB &&
if (hasOperation(ISD::USUBSAT, VT)) {
         if (OpLHS == LHS) {
           // Look for a general sub with unsigned saturation first.
           // x >= y ? x-y : 0 --> usubsat x, y
           // x >  y ? x-y : 0 --> usubsat x, y
           if ((SatCC == ISD::SETUGE || SatCC == ISD::SETUGT) &&
               Other.getOpcode() == ISD::SUB && OpRHS == CondRHS)
             return DAG.getNode(ISD::USUBSAT, DL, VT, OpLHS, OpRHS);
 
-          if (auto *OpRHSBV = dyn_cast<BuildVectorSDNode>(OpRHS)) {
-            if (isa<BuildVectorSDNode>(CondRHS)) {
+          if (OpRHS.getOpcode() == ISD::BUILD_VECTOR ||
+              OpRHS.getOpcode() == ISD::SPLAT_VECTOR) {
+            if (CondRHS.getOpcode() == ISD::BUILD_VECTOR ||
+                CondRHS.getOpcode() == ISD::SPLAT_VECTOR) {
               // If the RHS is a constant we have to reverse the const
               // canonicalization.
               // x > C-1 ? x+-C : 0 --> usubsat x, C
               auto MatchUSUBSAT = [](ConstantSDNode *Op, ConstantSDNode *Cond) {
                 return (!Op && !Cond) ||
                        (Op && Cond &&
                         Cond->getAPIntValue() == (-Op->getAPIntValue() - 1));
               };
               if (SatCC == ISD::SETUGT && Other.getOpcode() == ISD::ADD &&
                   ISD::matchBinaryPredicate(OpRHS, CondRHS, MatchUSUBSAT,
                                             /*AllowUndefs*/ true)) {
                 OpRHS = DAG.getNode(ISD::SUB, DL, VT,
                                     DAG.getConstant(0, DL, VT), OpRHS);
                 return DAG.getNode(ISD::USUBSAT, DL, VT, OpLHS, OpRHS);
               }
 
               // Another special case: If C was a sign bit, the sub has been
               // canonicalized into a xor.
               // FIXME: Would it be better to use computeKnownBits to determine
               //        whether it's safe to decanonicalize the xor?
               // x s< 0 ? x^C : 0 --> usubsat x, C
-              if (auto *OpRHSConst = OpRHSBV->getConstantSplatNode()) {
-                if (SatCC == ISD::SETLT && Other.getOpcode() == ISD::XOR &&
-                    ISD::isBuildVectorAllZeros(CondRHS.getNode()) &&
-                    OpRHSConst->getAPIntValue().isSignMask()) {
-                  // Note that we have to rebuild the RHS constant here to
-                  // ensure we don't rely on particular values of undef lanes.
-                  OpRHS = DAG.getConstant(OpRHSConst->getAPIntValue(), DL, VT);
-                  return DAG.getNode(ISD::USUBSAT, DL, VT, OpLHS, OpRHS);
-                }
-              }
-            }
-          }
-        }
-      }
-    }
-  }
+              APInt SplatValue;
+                  ISD::isConstantSplatVector(OpRHS.getNode(), SplatValue) &&
+                  ISD::isConstantSplatVectorAllZeros(CondRHS.getNode()) &&
+                  SplatValue.isSignMask()) {
+                OpRHS = DAG.getConstant(SplatValue, DL, VT);
 
   if (SimplifySelectOps(N, N1, N2))
     return SDValue(N, 0);  // Don't revisit N.
 
   // Fold (vselect all_ones, N1, N2) -> N1
   if (ISD::isConstantSplatVectorAllOnes(N0.getNode()))
     return N1;
   // Fold (vselect all_zeros, N1, N2) -> N2

llvm/test/CodeGen/RISCV/rvv/combine-sats.ll

; CHECK-NEXT:    ret
   %v2 = select <2 x i1> %cmp, <2 x i64> %v1, <2 x i64> zeroinitializer
   ret <2 x i64> %v2
 }
 
 define <vscale x 2 x i64> @vselect_sub_nxv2i64(<vscale x 2 x i64> %a0, <vscale x 2 x i64> %a1) {
 ; CHECK-LABEL: vselect_sub_nxv2i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e64, m2, ta, mu
-; CHECK-NEXT:    vmsleu.vv v0, v10, v8
-; CHECK-NEXT:    vsub.vv v26, v8, v10
-; CHECK-NEXT:    vmv.v.i v28, 0
-; CHECK-NEXT:    vmerge.vvm v8, v28, v26, v0
+; CHECK-NEXT:    vssubu.vv v8, v8, v10
 ; CHECK-NEXT:    ret
   %cmp = icmp uge <vscale x 2 x i64> %a0, %a1
   %v1 = sub <vscale x 2 x i64> %a0, %a1
   %v2 = select <vscale x 2 x i1> %cmp, <vscale x 2 x i64> %v1, <vscale x 2 x i64> zeroinitializer
   ret <vscale x 2 x i64> %v2
 }
 
 define <8 x i16> @vselect_sub_2_v8i16(<8 x i16> %x, i16 zeroext %w) nounwind {
entry:
   %res = select <8 x i1> %1, <8 x i16> zeroinitializer, <8 x i16> %2
   ret <8 x i16> %res
 }
 
 define <vscale x 8 x i16> @vselect_sub_2_nxv8i16(<vscale x 8 x i16> %x, i16 zeroext %w) nounwind {
 ; CHECK-LABEL: vselect_sub_2_nxv8i16:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    vsetvli a1, zero, e16, m2, ta, mu
-; CHECK-NEXT:    vmsltu.vx v0, v8, a0
-; CHECK-NEXT:    vsub.vx v26, v8, a0
-; CHECK-NEXT:    vmerge.vim v8, v26, 0, v0
+; CHECK-NEXT:    vssubu.vx v8, v8, a0
 ; CHECK-NEXT:    ret
 entry:
   %0 = insertelement <vscale x 8 x i16> undef, i16 %w, i32 0
   %broadcast15 = shufflevector <vscale x 8 x i16> %0, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
   %1 = icmp ult <vscale x 8 x i16> %x, %broadcast15
   %2 = sub <vscale x 8 x i16> %x, %broadcast15
   %res = select <vscale x 8 x i1> %1, <vscale x 8 x i16> zeroinitializer, <vscale x 8 x i16> %2
   ret <vscale x 8 x i16> %res
; CHECK-NEXT:    ret
   %cmp = icmp ugt <2 x i64> %a0, <i64 5, i64 5>
   %v2 = select <2 x i1> %cmp, <2 x i64> %v1, <2 x i64> zeroinitializer
   ret <2 x i64> %v2
 }
 
 define <vscale x 2 x i64> @vselect_add_const_nxv2i64(<vscale x 2 x i64> %a0) {
 ; CHECK-LABEL: vselect_add_const_nxv2i64:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetvli a0, zero, e64, m2, ta, mu
-; CHECK-NEXT:    vadd.vi v26, v8, -6
-; CHECK-NEXT:    vmsgtu.vi v0, v8, 5
-; CHECK-NEXT:    vmv.v.i v28, 0
-; CHECK-NEXT:    vmerge.vvm v8, v28, v26, v0
+; CHECK-NEXT:    addi a0, zero, 6
+; CHECK-NEXT:    vsetvli a1, zero, e64, m2, ta, mu
+; CHECK-NEXT:    vssubu.vx v8, v8, a0
 ; CHECK-NEXT:    ret
   %cm1 = insertelement <vscale x 2 x i64> poison, i64 -6, i32 0
   %splatcm1 = shufflevector <vscale x 2 x i64> %cm1, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %nc = insertelement <vscale x 2 x i64> poison, i64 5, i32 0
   %splatnc = shufflevector <vscale x 2 x i64> %nc, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %v1 = add <vscale x 2 x i64> %a0, %splatcm1
   %cmp = icmp ugt <vscale x 2 x i64> %a0, %splatnc
   %v2 = select <vscale x 2 x i1> %cmp, <vscale x 2 x i64> %v1, <vscale x 2 x i64> zeroinitializer
; RV64-NEXT:    ret
   %v2 = select <2 x i1> %cmp, <2 x i16> %v1, <2 x i16> zeroinitializer
   ret <2 x i16> %v2
 }
 
 define <vscale x 2 x i16> @vselect_add_const_signbit_nxv2i16(<vscale x 2 x i16> %a0) {
 ; RV32-LABEL: vselect_add_const_signbit_nxv2i16:
 ; RV32:       # %bb.0:
 ; RV32-NEXT:    lui a0, 8
-; RV32-NEXT:    addi a0, a0, -2
+; RV32-NEXT:    addi a0, a0, -1
 ; RV32-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
-; RV32-NEXT:    vmsgtu.vx v0, v8, a0
-; RV32-NEXT:    lui a0, 1048568
-; RV32-NEXT:    addi a0, a0, 1
-; RV32-NEXT:    vadd.vx v25, v8, a0
-; RV32-NEXT:    vmv.v.i v26, 0
-; RV32-NEXT:    vmerge.vvm v8, v26, v25, v0
+; RV32-NEXT:    vssubu.vx v8, v8, a0
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: vselect_add_const_signbit_nxv2i16:
 ; RV64:       # %bb.0:
 ; RV64-NEXT:    lui a0, 8
-; RV64-NEXT:    addiw a0, a0, -2
+; RV64-NEXT:    addiw a0, a0, -1
 ; RV64-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
-; RV64-NEXT:    vmsgtu.vx v0, v8, a0
-; RV64-NEXT:    lui a0, 1048568
-; RV64-NEXT:    addiw a0, a0, 1
-; RV64-NEXT:    vadd.vx v25, v8, a0
-; RV64-NEXT:    vmv.v.i v26, 0
-; RV64-NEXT:    vmerge.vvm v8, v26, v25, v0
+; RV64-NEXT:    vssubu.vx v8, v8, a0
 ; RV64-NEXT:    ret
   %cm1 = insertelement <vscale x 2 x i16> poison, i16 32766, i32 0
   %splatcm1 = shufflevector <vscale x 2 x i16> %cm1, <vscale x 2 x i16> poison, <vscale x 2 x i32> zeroinitializer
   %nc = insertelement <vscale x 2 x i16> poison, i16 -32767, i32 0
   %splatnc = shufflevector <vscale x 2 x i16> %nc, <vscale x 2 x i16> poison, <vscale x 2 x i32> zeroinitializer
   %cmp = icmp ugt <vscale x 2 x i16> %a0, %splatcm1
   %v1 = add <vscale x 2 x i16> %a0, %splatnc
   %v2 = select <vscale x 2 x i1> %cmp, <vscale x 2 x i16> %v1, <vscale x 2 x i16> zeroinitializer
; CHECK-NEXT:    ret
   %v1 = xor <2 x i16> %a0, <i16 -32768, i16 -32768>
   %v2 = select <2 x i1> %cmp, <2 x i16> %v1, <2 x i16> zeroinitializer
   ret <2 x i16> %v2
 }
 
 define <vscale x 2 x i16> @vselect_xor_const_signbit_nxv2i16(<vscale x 2 x i16> %a0) {
 ; CHECK-LABEL: vselect_xor_const_signbit_nxv2i16:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetvli a0, zero, e16, mf2, ta, mu
-; CHECK-NEXT:    vmsle.vi v0, v8, -1
-; CHECK-NEXT:    vmv.v.i v25, 0
-; CHECK-NEXT:    lui a0, 1048568
-; CHECK-NEXT:    vxor.vx v26, v8, a0
-; CHECK-NEXT:    vmerge.vvm v8, v25, v26, v0
+; CHECK-NEXT:    lui a0, 8
+; CHECK-NEXT:    vsetvli a1, zero, e16, mf2, ta, mu
+; CHECK-NEXT:    vssubu.vx v8, v8, a0
 ; CHECK-NEXT:    ret
   %cmp = icmp slt <vscale x 2 x i16> %a0, zeroinitializer
   %ins = insertelement <vscale x 2 x i16> poison, i16 -32768, i32 0
   %splat = shufflevector <vscale x 2 x i16> %ins, <vscale x 2 x i16> poison, <vscale x 2 x i32> zeroinitializer
   %v1 = xor <vscale x 2 x i16> %a0, %splat
   %v2 = select <vscale x 2 x i1> %cmp, <vscale x 2 x i16> %v1, <vscale x 2 x i16> zeroinitializer
   ret <vscale x 2 x i16> %v2
 }
; CHECK-NEXT:    ret
   %v2 = select <2 x i1> %cmp, <2 x i64> %v1, <2 x i64> <i64 -1, i64 -1>
   ret <2 x i64> %v2
 }
 
 define <vscale x 2 x i64> @vselect_add_nxv2i64(<vscale x 2 x i64> %a0, <vscale x 2 x i64> %a1) {
 ; CHECK-LABEL: vselect_add_nxv2i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e64, m2, ta, mu
-; CHECK-NEXT:    vadd.vv v26, v8, v10
-; CHECK-NEXT:    vmsleu.vv v0, v8, v26
-; CHECK-NEXT:    vmv.v.i v28, -1
-; CHECK-NEXT:    vmerge.vvm v8, v28, v26, v0
+; CHECK-NEXT:    vsaddu.vv v8, v8, v10
 ; CHECK-NEXT:    ret
   %v1 = add <vscale x 2 x i64> %a0, %a1
   %cmp = icmp ule <vscale x 2 x i64> %a0, %v1
   %allones = insertelement <vscale x 2 x i64> poison, i64 -1, i32 0
   %splatallones = shufflevector <vscale x 2 x i64> %allones, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %v2 = select <vscale x 2 x i1> %cmp, <vscale x 2 x i64> %v1, <vscale x 2 x i64> %splatallones
   ret <vscale x 2 x i64> %v2
 }
; CHECK-NEXT:    ret
   %v2 = select <2 x i1> %cmp, <2 x i64> %v1, <2 x i64> <i64 -1, i64 -1>
   ret <2 x i64> %v2
 }
 
 define <vscale x 2 x i64> @vselect_add_const_2_nxv2i64(<vscale x 2 x i64> %a0) {
 ; CHECK-LABEL: vselect_add_const_2_nxv2i64:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e64, m2, ta, mu
-; CHECK-NEXT:    vadd.vi v26, v8, 6
-; CHECK-NEXT:    vmsleu.vi v0, v8, -7
-; CHECK-NEXT:    vmv.v.i v28, -1
-; CHECK-NEXT:    vmerge.vvm v8, v28, v26, v0
+; CHECK-NEXT:    vsaddu.vi v8, v8, 6
 ; CHECK-NEXT:    ret
   %cm1 = insertelement <vscale x 2 x i64> poison, i64 6, i32 0
   %splatcm1 = shufflevector <vscale x 2 x i64> %cm1, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %nc = insertelement <vscale x 2 x i64> poison, i64 -7, i32 0
   %splatnc = shufflevector <vscale x 2 x i64> %nc, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %v1 = add <vscale x 2 x i64> %a0, %splatcm1
   %cmp = icmp ule <vscale x 2 x i64> %a0, %splatnc
   %allones = insertelement <vscale x 2 x i64> poison, i64 -1, i32 0