[DAGCombine] Improve (sra (sra x, c1), c2) -> (sra x, (add c1, c2)) folding

[RKSimon]

Summary

Add support for cases where only some c1+c2 results exceed the max bitshift, clamping accordingly.

[andrew.zhogin]

Isn't it a little regression for scalar types?

N1.getValueType() == N0.getOperand(1).getValueType()

Instead of

zeroExtendToMatch(c1, c2, 1 /* Overflow Bit */);

There is no vector analog of zeroExtendToMatch currently, right?

Btw, I don't think I have rights to accept the revision.

[RKSimon]

Resurrecting this old patch, refactored to correctly handle overflows of the sum of shifts.

[spatel]

test/CodeGen/X86/combine-sra.ll

144–145

To provide better coverage, add (or adjust this) test where both of the component shifts are below bitwidth, but the sum exceeds bitwidth?

[RKSimon]

test/CodeGen/X86/combine-sra.ll

144–145

See @combine_vec_ashr_ashr2 above

[spatel]

LGTM. I notice we don't fold the non-splat case of this pattern in IR. Do you think it's worth adding there too, or backend is good enough?

[RKSimon]

In D35722#1202997, @spatel wrote:

LGTM. I notice we don't fold the non-splat case of this pattern in IR. Do you think it's worth adding there too, or backend is good enough?

Add it if you can, properly supporting vectors is always a good idea ;-)

Is there the equivalent logical shifts -> zero combine?

[spatel]

In D35722#1203007, @RKSimon wrote:

In D35722#1202997, @spatel wrote:

LGTM. I notice we don't fold the non-splat case of this pattern in IR. Do you think it's worth adding there too, or backend is good enough?

Add it if you can, properly supporting vectors is always a good idea ;-)

Is there the equivalent logical shifts -> zero combine?

AFAIK, no. I think splats are well supported now in instcombine, but we haven't generalized many transforms for non-splats.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (VT.isVector())
                                ExtVT, VT.getVectorNumElements());
     if ((!LegalOperations ||
          TLI.isOperationLegal(ISD::SIGN_EXTEND_INREG, ExtVT)))
       return DAG.getNode(ISD::SIGN_EXTEND_INREG, SDLoc(N), VT,
                          N0.getOperand(0), DAG.getValueType(ExtVT));
   }
 
   // fold (sra (sra x, c1), c2) -> (sra x, (add c1, c2))
-  if (N0.getOpcode() == ISD::SRA) {
+  // clamp (add c1, c2) to max shift.
+  if (N0.getOpcode() == ISD::SRA &&
+      isConstantOrConstantVector(N1, /* NoOpaques */ true) &&
+      isConstantOrConstantVector(N0.getOperand(1), /* NoOpaques */ true) &&
+      N1.getValueType() == N0.getOperand(1).getValueType()) {
     SDLoc DL(N);
     EVT ShiftVT = N1.getValueType();
-
-    auto MatchOutOfRange = [OpSizeInBits](ConstantSDNode *LHS,
-                                          ConstantSDNode *RHS) {
-      APInt c1 = LHS->getAPIntValue();
-      APInt c2 = RHS->getAPIntValue();
-      zeroExtendToMatch(c1, c2, 1 /* Overflow Bit */);
-      return (c1 + c2).uge(OpSizeInBits);
-    };
-    if (matchBinaryPredicate(N1, N0.getOperand(1), MatchOutOfRange))
-      return DAG.getNode(ISD::SRA, DL, VT, N0.getOperand(0),
-                         DAG.getConstant(OpSizeInBits - 1, DL, ShiftVT));
-
-    auto MatchInRange = [OpSizeInBits](ConstantSDNode *LHS,
-                                       ConstantSDNode *RHS) {
-      APInt c1 = LHS->getAPIntValue();
-      APInt c2 = RHS->getAPIntValue();
-      zeroExtendToMatch(c1, c2, 1 /* Overflow Bit */);
-      return (c1 + c2).ult(OpSizeInBits);
-    };
-    if (matchBinaryPredicate(N1, N0.getOperand(1), MatchInRange)) {
-      SDValue Sum = DAG.getNode(ISD::ADD, DL, ShiftVT, N1, N0.getOperand(1));
-      return DAG.getNode(ISD::SRA, DL, VT, N0.getOperand(0), Sum);
-    }
+    SDValue Limit = DAG.getConstant(OpSizeInBits - 1, DL, ShiftVT);
+    SDValue Clamp = DAG.getNode(ISD::UMIN, DL, ShiftVT, Sum, Limit);
+    return DAG.getNode(ISD::SRA, DL, VT, N0.getOperand(0), Clamp);
   }
 
   // fold (sra (shl X, m), (sub result_size, n))
   // -> (sign_extend (trunc (shl X, (sub (sub result_size, n), m)))) for
   // result_size - n != m.
   // If truncate is free for the target sext(shl) is likely to result in better
   // code.
   if (N0.getOpcode() == ISD::SHL && N1C) {

test/CodeGen/X86/combine-sra.ll

 }
 
 define <4 x i32> @combine_vec_ashr_ashr3(<4 x i32> %x) {
 ; SSE-LABEL: combine_vec_ashr_ashr3:
 ; SSE:       # BB#0:
 ; SSE-NEXT:    movdqa %xmm0, %xmm1
 ; SSE-NEXT:    psrad $27, %xmm1
 ; SSE-NEXT:    movdqa %xmm0, %xmm2
-; SSE-NEXT:    psrad $5, %xmm2
+; SSE-NEXT:    psrad $15, %xmm2
 ; SSE-NEXT:    pblendw {{.*#+}} xmm2 = xmm2[0,1,2,3],xmm1[4,5,6,7]
-; SSE-NEXT:    movdqa %xmm0, %xmm1
-; SSE-NEXT:    psrad $31, %xmm1
-; SSE-NEXT:    psrad $1, %xmm0
-; SSE-NEXT:    pblendw {{.*#+}} xmm0 = xmm0[0,1,2,3],xmm1[4,5,6,7]
-; SSE-NEXT:    pblendw {{.*#+}} xmm0 = xmm0[0,1],xmm2[2,3],xmm0[4,5],xmm2[6,7]
-; SSE-NEXT:    movdqa %xmm0, %xmm1
-; SSE-NEXT:    psrad $10, %xmm1
-; SSE-NEXT:    pblendw {{.*#+}} xmm1 = xmm1[0,1,2,3],xmm0[4,5,6,7]
 ; SSE-NEXT:    psrad $31, %xmm0
-; SSE-NEXT:    pblendw {{.*#+}} xmm0 = xmm0[0,1],xmm1[2,3],xmm0[4,5],xmm1[6,7]
+; SSE-NEXT:    pblendw {{.*#+}} xmm0 = xmm0[0,1],xmm2[2,3],xmm0[4,5],xmm2[6,7]
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: combine_vec_ashr_ashr3:
 ; AVX:       # BB#0:
 ; AVX-NEXT:    vpsravd {{.*}}(%rip), %xmm0, %xmm0
-; AVX-NEXT:    vpsravd {{.*}}(%rip), %xmm0, %xmm0
 ; AVX-NEXT:    retq
   %1 = ashr <4 x i32> %x, <i32  1, i32  5, i32 50, i32 27>
   %2 = ashr <4 x i32> %1, <i32 33, i32 10, i32 33, i32  0>
   ret <4 x i32> %2
 }
 
 ; fold (sra x, (trunc (and y, c))) -> (sra x, (and (trunc y), (trunc c))).
 define <4 x i32> @combine_vec_ashr_trunc_and(<4 x i32> %x, <4 x i64> %y) {
 ; SSE-LABEL: combine_vec_ashr_trunc_and:
 ; SSE:       # BB#0:
 ; SSE-NEXT:    shufps {{.*#+}} xmm1 = xmm1[0,2],xmm2[0,2]