[SDAG] enable binop identity constant folds for shifts

[RKSimon]

Summary

Add shl/srl/sra to the list of ops that we canonicalize with a select to expose an identity merge

[pengfei]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2167

Can it apply to ROTL/ROTR ?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2167

Yes, we'd have to account for the implicit modulo though.

There's plenty of other ops we could attempt with this as well, including some non-binops (funnel shifts and fma for instance) - I'm also wondering what target binops will need to be supported.

[LuoYuanke]

LGTM, thanks.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (ConstantFPSDNode *C = isConstOrConstSplatFP(V)) {
       case ISD::FMUL: // X * 1.0 --> X
       case ISD::FDIV: // X / 1.0 --> X
         return C->isExactlyValue(1.0);
       }
     }
     if (ConstantSDNode *C = isConstOrConstSplat(V)) {
       switch (Opcode) {
       case ISD::SUB: // X - 0 --> X
+      case ISD::SHL: // X << 0 --> X
+      case ISD::SRA: // X s>> 0 --> X
+      case ISD::SRL: // X u>> 0 --> X
         return C->isZero();
       }
     }
     return false;
   };
 
   // This transform increases uses of N0, so freeze it to be safe.
   // binop N0, (vselect Cond, IDC, FVal) --> vselect Cond, N0, (binop N0, FVal)

llvm/test/CodeGen/X86/vector-bo-select.ll

 ; AVX512F-NEXT:    vpsllvd %xmm0, %xmm1, %xmm0
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: shl_v4i32:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    vpslld $31, %xmm0, %xmm0
 ; AVX512VL-NEXT:    vptestmd %xmm0, %xmm0, %k1
-; AVX512VL-NEXT:    vmovdqa32 %xmm2, %xmm0 {%k1} {z}
-; AVX512VL-NEXT:    vpsllvd %xmm0, %xmm1, %xmm0
+; AVX512VL-NEXT:    vpsllvd %xmm2, %xmm1, %xmm1 {%k1}
+; AVX512VL-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX512VL-NEXT:    retq
   %s = select <4 x i1> %b, <4 x i32> %y, <4 x i32> zeroinitializer
   %r = shl <4 x i32> %x, %s
   ret <4 x i32> %r
 }
 
 ; negative test - shl is not commutative; there is no identity constant for operand 0
 

 ; AVX2-NEXT:    vpsllvd %ymm0, %ymm1, %ymm0
 ; AVX2-NEXT:    vpsllvd %ymm4, %ymm2, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: shl_v16i32_swap:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    vpmovsxbd %xmm0, %zmm0
 ; AVX512-NEXT:    vpslld $31, %zmm0, %zmm0
-; AVX512-NEXT:    vptestnmd %zmm0, %zmm0, %k1
-; AVX512-NEXT:    vmovdqa32 %zmm2, %zmm0 {%k1} {z}
-; AVX512-NEXT:    vpsllvd %zmm0, %zmm1, %zmm0
+; AVX512-NEXT:    vptestmd %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpsllvd %zmm2, %zmm1, %zmm0
+; AVX512-NEXT:    vmovdqa32 %zmm1, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %s = select <16 x i1> %b, <16 x i32> zeroinitializer, <16 x i32> %y
   %r = shl <16 x i32> %x, %s
   ret <16 x i32> %r
 }
 
 ; negative test - shl is not commutative; there is no identity constant for operand 0
 

 ; AVX512F-NEXT:    kmovw %edi, %k1
 ; AVX512F-NEXT:    vmovdqa32 %zmm1, %zmm1 {%k1} {z}
 ; AVX512F-NEXT:    vpsllvd %ymm1, %ymm0, %ymm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: shl_v8i32_cast_cond:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    kmovw %edi, %k1
-; AVX512VL-NEXT:    vmovdqa32 %ymm1, %ymm1 {%k1} {z}
-; AVX512VL-NEXT:    vpsllvd %ymm1, %ymm0, %ymm0
+; AVX512VL-NEXT:    vpsllvd %ymm1, %ymm0, %ymm0 {%k1}
 ; AVX512VL-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i32> %y, <8 x i32> zeroinitializer
   %r = shl <8 x i32> %x, %s
   ret <8 x i32> %r
 }
 
 define <8 x i64> @shl_v8i64_cast_cond(i8 noundef zeroext %pb, <8 x i64> noundef %x, <8 x i64> noundef %y) {

 ; AVX2-NEXT:    vpand %ymm2, %ymm4, %ymm2
 ; AVX2-NEXT:    vpsllvq %ymm2, %ymm0, %ymm0
 ; AVX2-NEXT:    vpsllvq %ymm3, %ymm1, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: shl_v8i64_cast_cond:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    kmovw %edi, %k1
-; AVX512-NEXT:    vmovdqa64 %zmm1, %zmm1 {%k1} {z}
-; AVX512-NEXT:    vpsllvq %zmm1, %zmm0, %zmm0
+; AVX512-NEXT:    vpsllvq %zmm1, %zmm0, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i64> %y, <8 x i64> zeroinitializer
   %r = shl <8 x i64> %x, %s
   ret <8 x i64> %r
 }
 
 define <4 x i32> @lshr_v4i32(<4 x i1> %b, <4 x i32> noundef %x, <4 x i32> noundef %y) {

 ; AVX512F-NEXT:    vpsrlvd %xmm0, %xmm1, %xmm0
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: lshr_v4i32:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    vpslld $31, %xmm0, %xmm0
 ; AVX512VL-NEXT:    vptestmd %xmm0, %xmm0, %k1
-; AVX512VL-NEXT:    vmovdqa32 %xmm2, %xmm0 {%k1} {z}
-; AVX512VL-NEXT:    vpsrlvd %xmm0, %xmm1, %xmm0
+; AVX512VL-NEXT:    vpsrlvd %xmm2, %xmm1, %xmm1 {%k1}
+; AVX512VL-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX512VL-NEXT:    retq
   %s = select <4 x i1> %b, <4 x i32> %y, <4 x i32> zeroinitializer
   %r = lshr <4 x i32> %x, %s
   ret <4 x i32> %r
 }
 
 ; negative test - lshr is not commutative; there is no identity constant for operand 0
 

 ; AVX2-NEXT:    vpsrlvd %ymm0, %ymm1, %ymm0
 ; AVX2-NEXT:    vpsrlvd %ymm4, %ymm2, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: lshr_v16i32_swap:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    vpmovsxbd %xmm0, %zmm0
 ; AVX512-NEXT:    vpslld $31, %zmm0, %zmm0
-; AVX512-NEXT:    vptestnmd %zmm0, %zmm0, %k1
-; AVX512-NEXT:    vmovdqa32 %zmm2, %zmm0 {%k1} {z}
-; AVX512-NEXT:    vpsrlvd %zmm0, %zmm1, %zmm0
+; AVX512-NEXT:    vptestmd %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpsrlvd %zmm2, %zmm1, %zmm0
+; AVX512-NEXT:    vmovdqa32 %zmm1, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %s = select <16 x i1> %b, <16 x i32> zeroinitializer, <16 x i32> %y
   %r = lshr <16 x i32> %x, %s
   ret <16 x i32> %r
 }
 
 ; negative test - lshr is not commutative; there is no identity constant for operand 0
 

 ; AVX512F-NEXT:    kmovw %edi, %k1
 ; AVX512F-NEXT:    vmovdqa32 %zmm1, %zmm1 {%k1} {z}
 ; AVX512F-NEXT:    vpsrlvd %ymm1, %ymm0, %ymm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: lshr_v8i32_cast_cond:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    kmovw %edi, %k1
-; AVX512VL-NEXT:    vmovdqa32 %ymm1, %ymm1 {%k1} {z}
-; AVX512VL-NEXT:    vpsrlvd %ymm1, %ymm0, %ymm0
+; AVX512VL-NEXT:    vpsrlvd %ymm1, %ymm0, %ymm0 {%k1}
 ; AVX512VL-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i32> %y, <8 x i32> zeroinitializer
   %r = lshr <8 x i32> %x, %s
   ret <8 x i32> %r
 }
 
 define <8 x i64> @lshr_v8i64_cast_cond(i8 noundef zeroext %pb, <8 x i64> noundef %x, <8 x i64> noundef %y) {

 ; AVX2-NEXT:    vpand %ymm2, %ymm4, %ymm2
 ; AVX2-NEXT:    vpsrlvq %ymm2, %ymm0, %ymm0
 ; AVX2-NEXT:    vpsrlvq %ymm3, %ymm1, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: lshr_v8i64_cast_cond:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    kmovw %edi, %k1
-; AVX512-NEXT:    vmovdqa64 %zmm1, %zmm1 {%k1} {z}
-; AVX512-NEXT:    vpsrlvq %zmm1, %zmm0, %zmm0
+; AVX512-NEXT:    vpsrlvq %zmm1, %zmm0, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i64> %y, <8 x i64> zeroinitializer
   %r = lshr <8 x i64> %x, %s
   ret <8 x i64> %r
 }
 
 define <4 x i32> @ashr_v4i32(<4 x i1> %b, <4 x i32> noundef %x, <4 x i32> noundef %y) {

 ; AVX512F-NEXT:    vpsravd %xmm0, %xmm1, %xmm0
 ; AVX512F-NEXT:    vzeroupper
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: ashr_v4i32:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    vpslld $31, %xmm0, %xmm0
 ; AVX512VL-NEXT:    vptestmd %xmm0, %xmm0, %k1
-; AVX512VL-NEXT:    vmovdqa32 %xmm2, %xmm0 {%k1} {z}
-; AVX512VL-NEXT:    vpsravd %xmm0, %xmm1, %xmm0
+; AVX512VL-NEXT:    vpsravd %xmm2, %xmm1, %xmm1 {%k1}
+; AVX512VL-NEXT:    vmovdqa %xmm1, %xmm0
 ; AVX512VL-NEXT:    retq
   %s = select <4 x i1> %b, <4 x i32> %y, <4 x i32> zeroinitializer
   %r = ashr <4 x i32> %x, %s
   ret <4 x i32> %r
 }
 
 ; negative test - ashr is not commutative; there is no identity constant for operand 0
 

 ; AVX2-NEXT:    vpsravd %ymm0, %ymm1, %ymm0
 ; AVX2-NEXT:    vpsravd %ymm4, %ymm2, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: ashr_v16i32_swap:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    vpmovsxbd %xmm0, %zmm0
 ; AVX512-NEXT:    vpslld $31, %zmm0, %zmm0
-; AVX512-NEXT:    vptestnmd %zmm0, %zmm0, %k1
-; AVX512-NEXT:    vmovdqa32 %zmm2, %zmm0 {%k1} {z}
-; AVX512-NEXT:    vpsravd %zmm0, %zmm1, %zmm0
+; AVX512-NEXT:    vptestmd %zmm0, %zmm0, %k1
+; AVX512-NEXT:    vpsravd %zmm2, %zmm1, %zmm0
+; AVX512-NEXT:    vmovdqa32 %zmm1, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %s = select <16 x i1> %b, <16 x i32> zeroinitializer, <16 x i32> %y
   %r = ashr <16 x i32> %x, %s
   ret <16 x i32> %r
 }
 
 ; negative test - ashr is not commutative; there is no identity constant for operand 0
 

 ; AVX512F-NEXT:    kmovw %edi, %k1
 ; AVX512F-NEXT:    vmovdqa32 %zmm1, %zmm1 {%k1} {z}
 ; AVX512F-NEXT:    vpsravd %ymm1, %ymm0, %ymm0
 ; AVX512F-NEXT:    retq
 ;
 ; AVX512VL-LABEL: ashr_v8i32_cast_cond:
 ; AVX512VL:       # %bb.0:
 ; AVX512VL-NEXT:    kmovw %edi, %k1
-; AVX512VL-NEXT:    vmovdqa32 %ymm1, %ymm1 {%k1} {z}
-; AVX512VL-NEXT:    vpsravd %ymm1, %ymm0, %ymm0
+; AVX512VL-NEXT:    vpsravd %ymm1, %ymm0, %ymm0 {%k1}
 ; AVX512VL-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i32> %y, <8 x i32> zeroinitializer
   %r = ashr <8 x i32> %x, %s
   ret <8 x i32> %r
 }
 
 define <8 x i64> @ashr_v8i64_cast_cond(i8 noundef zeroext %pb, <8 x i64> noundef %x, <8 x i64> noundef %y) {

 ; AVX2-NEXT:    vpsrlvq %ymm3, %ymm1, %ymm1
 ; AVX2-NEXT:    vpxor %ymm2, %ymm1, %ymm1
 ; AVX2-NEXT:    vpsubq %ymm2, %ymm1, %ymm1
 ; AVX2-NEXT:    retq
 ;
 ; AVX512-LABEL: ashr_v8i64_cast_cond:
 ; AVX512:       # %bb.0:
 ; AVX512-NEXT:    kmovw %edi, %k1
-; AVX512-NEXT:    vmovdqa64 %zmm1, %zmm1 {%k1} {z}
-; AVX512-NEXT:    vpsravq %zmm1, %zmm0, %zmm0
+; AVX512-NEXT:    vpsravq %zmm1, %zmm0, %zmm0 {%k1}
 ; AVX512-NEXT:    retq
   %b = bitcast i8 %pb to <8 x i1>
   %s = select <8 x i1> %b, <8 x i64> %y, <8 x i64> zeroinitializer
   %r = ashr <8 x i64> %x, %s
   ret <8 x i64> %r
 }