[InstCombine][SSE] Add DemandedElts support for PACKSS/PACKUS instructions

[RKSimon]

Summary

Simplify a packss/packus truncation based on the elements of the mask that are actually demanded.

[mkuper]

lib/Transforms/InstCombine/InstCombineCalls.cpp

1831 ↗

(On Diff #84576)

This isn't a question about this patch (and maybe I'm not the person who should be reviewing it :-) ), but I'm now confused about how this works at all.

What's the benefit of calling SimplifyDemandedVectorElts() on an instruction when you're demanding all elements, the instruction you're calling this on will also demand all the elements of its operands, and throwing UndefElts away? Is this just so that things have a chance to happen in the middle of the tree?

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1475

Add a TODO for the AVX512 versions?

1495

Can you make this more verbose? It's a bit opaque unless you already know what the layout looks like.

1497

i -> Lane, j -> Elt (or something similar)?

1527

Maybe make UndefElts an array of 2 APInts? This will make this code slightly prettier.

[craig.topper]

lib/Transforms/InstCombine/InstCombineCalls.cpp

1831 ↗

(On Diff #84576)

Yeah this probably has no value for packs. Its useful for things like scalar intrinsics because they demand less elements from their inputs no matter what future users demand from the intrinsic.

I think this means PSHUFB doesn't need code in InstCombineCalls either.

[RKSimon]

I've removed the demanded elts handling of pshufb/vpermil from InstCombiner::visitCallInst

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1527

UndefElts is a pass by reference argument that is used by the caller so I can't do much to change it. I agree that the code is messy but it is probably as good as it can be.

[RKSimon]

Removed unnecessary handling in InstCombiner::visitCallInst, cleaned up code as requested.

[mkuper]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1527

Sorry, I meant "make UndefElts2 and UndefElts3 a 2-element array".
I think having this code with a nested loop would be nicer than duplicating each line twice.

[RKSimon]

Removed duplication of each operand by putting the demanded elts code in a loop.

[mkuper]

One more minor naming nit, other than that, LGTM.

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1496

i -> OpNum?

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

case Intrinsic::x86_avx512_pmulu_dq_512: {
       UndefElts3 = APInt(InnerVWidth, 0);
       TmpV = SimplifyDemandedVectorElts(Op1, InnerDemandedElts, UndefElts3,
                                         Depth + 1);
       if (TmpV) { II->setArgOperand(1, TmpV); MadeChange = true; }
 
       break;
     }
 
+    case Intrinsic::x86_sse2_packssdw_128:
+    case Intrinsic::x86_sse2_packsswb_128:
+    case Intrinsic::x86_sse2_packuswb_128:
+    case Intrinsic::x86_sse41_packusdw:
+    case Intrinsic::x86_avx2_packssdw:
+    case Intrinsic::x86_avx2_packsswb:
+    case Intrinsic::x86_avx2_packusdw:
+    case Intrinsic::x86_avx2_packuswb: {
+      // TODO Add support for Intrinsic::x86_avx512_mask_pack*
+      auto *Op0 = II->getArgOperand(0);
+      auto *Op1 = II->getArgOperand(1);
+      auto *Ty0 = Op0->getType();
+      unsigned InnerVWidth = Ty0->getVectorNumElements();
+      assert(VWidth == (InnerVWidth * 2) && "Unexpected input size");
+
+      APInt Op0DemandedElts(InnerVWidth, 0);
+      APInt Op1DemandedElts(InnerVWidth, 0);
+      unsigned NumLanes = Ty0->getPrimitiveSizeInBits() / 128;
+      unsigned VWidthPerLane = VWidth / NumLanes;
+      unsigned InnerVWidthPerLane = InnerVWidth / NumLanes;
+
+      // Per lane, pack the elements of the first input and then the second.
+      // e.g.
+      // v8i16 PACK(v4i32 X, v4i32 Y) - (X[0..3],Y[0..3])
+      // v32i8 PACK(v16i16 X, v16i16 Y) - (X[0..7],Y[0..7]),(X[8..15],Y[8..15])
+      for (unsigned Lane = 0; Lane != NumLanes; ++Lane)
+        for (unsigned Elt = 0; Elt != InnerVWidthPerLane; ++Elt) {
+          if (DemandedElts[(Lane * VWidthPerLane) + Elt])
+            Op0DemandedElts.setBit((Lane * InnerVWidthPerLane) + Elt);
+          if (DemandedElts[(Lane * VWidthPerLane) + Elt + InnerVWidthPerLane])
+            Op1DemandedElts.setBit((Lane * InnerVWidthPerLane) + Elt);
+        }
+
+      UndefElts2 = APInt(InnerVWidth, 0);
+      TmpV = SimplifyDemandedVectorElts(Op0, Op0DemandedElts, UndefElts2,
+                                        Depth + 1);
+      if (TmpV) {
+        II->setArgOperand(0, TmpV);
+        MadeChange = true;
+      }
+
+      UndefElts3 = APInt(InnerVWidth, 0);
+      TmpV = SimplifyDemandedVectorElts(Op1, Op1DemandedElts, UndefElts3,
+                                        Depth + 1);
+      if (TmpV) {
+        II->setArgOperand(1, TmpV);
+        MadeChange = true;
+      }
+
+      // Pack UNDEFs from the 2 inputs, one lane at a time.
+      UndefElts = APInt(VWidth, 0);
+      UndefElts2 = UndefElts2.zext(VWidth);
+      UndefElts3 = UndefElts3.zext(VWidth);
+      for (unsigned Lane = 0; Lane != NumLanes; ++Lane) {
+        APInt LaneElts2 = UndefElts2.lshr(InnerVWidthPerLane * Lane);
+        APInt LaneElts3 = UndefElts3.lshr(InnerVWidthPerLane * Lane);
+        LaneElts2 = LaneElts2.getLoBits(InnerVWidthPerLane);
+        LaneElts3 = LaneElts3.getLoBits(InnerVWidthPerLane);
+        LaneElts2 = LaneElts2.shl(InnerVWidthPerLane * (2 * Lane + 0));
+        LaneElts3 = LaneElts3.shl(InnerVWidthPerLane * (2 * Lane + 1));
+        UndefElts |= (LaneElts2 | LaneElts3);
+      }
+      break;
+    }
+
     // PSHUFB
     case Intrinsic::x86_ssse3_pshuf_b_128:
     case Intrinsic::x86_avx2_pshuf_b:
     case Intrinsic::x86_avx512_pshuf_b_512:
     // PERMILVAR
     case Intrinsic::x86_avx_vpermilvar_ps:
     case Intrinsic::x86_avx_vpermilvar_ps_256:
     case Intrinsic::x86_avx512_vpermilvar_ps_512:

test/Transforms/InstCombine/x86-pack.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ;
 ; Demanded Elts
 ;
 
 define <8 x i16> @elts_packssdw_128(<4 x i32> %a0, <4 x i32> %a1) {
 ; CHECK-LABEL: @elts_packssdw_128(
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> %a0, <4 x i32> undef, <4 x i32> <i32 3, i32 1, i32 undef, i32 undef>
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i32> %a1, <4 x i32> undef, <4 x i32> <i32 undef, i32 2, i32 1, i32 undef>
-; CHECK-NEXT:    [[TMP3:%.*]] = call <8 x i16> @llvm.x86.sse2.packssdw.128(<4 x i32> [[TMP1]], <4 x i32> [[TMP2]])
-; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <8 x i16> [[TMP3]], <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 7, i32 7, i32 7, i32 7>
-; CHECK-NEXT:    ret <8 x i16> [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x i16> @llvm.x86.sse2.packssdw.128(<4 x i32> %a0, <4 x i32> undef)
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <8 x i16> [[TMP1]], <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    ret <8 x i16> [[TMP2]]
 ;
   %1 = shufflevector <4 x i32> %a0, <4 x i32> undef, <4 x i32> <i32 3, i32 1, i32 undef, i32 undef>
   %2 = shufflevector <4 x i32> %a1, <4 x i32> undef, <4 x i32> <i32 undef, i32 2, i32 1, i32 undef>
   %3 = call <8 x i16> @llvm.x86.sse2.packssdw.128(<4 x i32> %1, <4 x i32> %2)
   %4 = shufflevector <8 x i16> %3, <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 7, i32 7, i32 7, i32 7>
   ret <8 x i16> %4
 }
 
 define <8 x i16> @elts_packusdw_128(<4 x i32> %a0, <4 x i32> %a1) {
 ; CHECK-LABEL: @elts_packusdw_128(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i32> %a0, i32 0, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x i32> %a1, i32 0, i32 3
-; CHECK-NEXT:    [[TMP3:%.*]] = call <8 x i16> @llvm.x86.sse41.packusdw(<4 x i32> [[TMP1]], <4 x i32> [[TMP2]])
-; CHECK-NEXT:    ret <8 x i16> [[TMP3]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x i16> @llvm.x86.sse41.packusdw(<4 x i32> %a0, <4 x i32> %a1)
+; CHECK-NEXT:    ret <8 x i16> [[TMP1]]
 ;
   %1 = insertelement <4 x i32> %a0, i32 0, i32 0
   %2 = insertelement <4 x i32> %a1, i32 0, i32 3
   %3 = call <8 x i16> @llvm.x86.sse41.packusdw(<4 x i32> %1, <4 x i32> %2)
   %4 = shufflevector <8 x i16> %3, <8 x i16> undef, <8 x i32> <i32 undef, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 undef>
   ret <8 x i16> %4
 }
 
 define <16 x i8> @elts_packsswb_128(<8 x i16> %a0, <8 x i16> %a1) {
 ; CHECK-LABEL: @elts_packsswb_128(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <8 x i16> %a0, i16 0, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <8 x i16> %a1, i16 0, i32 0
-; CHECK-NEXT:    [[TMP3:%.*]] = call <16 x i8> @llvm.x86.sse2.packsswb.128(<8 x i16> [[TMP1]], <8 x i16> [[TMP2]])
-; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <16 x i8> [[TMP3]], <16 x i8> undef, <16 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8>
-; CHECK-NEXT:    ret <16 x i8> [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <16 x i8> @llvm.x86.sse2.packsswb.128(<8 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>, <8 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>)
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <16 x i8> [[TMP1]], <16 x i8> undef, <16 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8>
+; CHECK-NEXT:    ret <16 x i8> [[TMP2]]
 ;
   %1 = insertelement <8 x i16> %a0, i16 0, i32 0
   %2 = insertelement <8 x i16> %a1, i16 0, i32 0
   %3 = call <16 x i8> @llvm.x86.sse2.packsswb.128(<8 x i16> %1, <8 x i16> %2)
   %4 = shufflevector <16 x i8> %3, <16 x i8> undef, <16 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8, i32 8>
   ret <16 x i8> %4
 }
 
 define <16 x i8> @elts_packuswb_128(<8 x i16> %a0, <8 x i16> %a1) {
 ; CHECK-LABEL: @elts_packuswb_128(
-; CHECK-NEXT:    [[TMP1:%.*]] = call <16 x i8> @llvm.x86.sse2.packuswb.128(<8 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>, <8 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>)
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <16 x i8> [[TMP1]], <16 x i8> undef, <16 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15>
-; CHECK-NEXT:    ret <16 x i8> [[TMP2]]
+; CHECK-NEXT:    ret <16 x i8> undef
 ;
   %1 = insertelement <8 x i16> undef, i16 0, i32 0
   %2 = insertelement <8 x i16> undef, i16 0, i32 0
   %3 = call <16 x i8> @llvm.x86.sse2.packuswb.128(<8 x i16> %1, <8 x i16> %2)
   %4 = shufflevector <16 x i8> %3, <16 x i8> undef, <16 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15, i32 15>
   ret <16 x i8> %4
 }
 
 define <16 x i16> @elts_packssdw_256(<8 x i32> %a0, <8 x i32> %a1) {
 ; CHECK-LABEL: @elts_packssdw_256(
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <8 x i32> %a0, <8 x i32> undef, <8 x i32> <i32 1, i32 0, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <8 x i32> %a1, <8 x i32> undef, <8 x i32> <i32 undef, i32 2, i32 1, i32 undef, i32 undef, i32 6, i32 5, i32 undef>
-; CHECK-NEXT:    [[TMP3:%.*]] = call <16 x i16> @llvm.x86.avx2.packssdw(<8 x i32> [[TMP1]], <8 x i32> [[TMP2]])
-; CHECK-NEXT:    ret <16 x i16> [[TMP3]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <16 x i16> @llvm.x86.avx2.packssdw(<8 x i32> %a0, <8 x i32> undef)
+; CHECK-NEXT:    ret <16 x i16> [[TMP1]]
 ;
   %1 = shufflevector <8 x i32> %a0, <8 x i32> undef, <8 x i32> <i32 1, i32 0, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   %2 = shufflevector <8 x i32> %a1, <8 x i32> undef, <8 x i32> <i32 undef, i32 2, i32 1, i32 undef, i32 undef, i32 6, i32 5, i32 undef>
   %3 = call <16 x i16> @llvm.x86.avx2.packssdw(<8 x i32> %1, <8 x i32> %2)
   %4 = shufflevector <16 x i16> %3, <16 x i16> undef, <16 x i32> <i32 undef, i32 undef, i32 2, i32 3, i32 4, i32 undef, i32 undef, i32 7, i32 8, i32 undef, i32 undef, i32 11, i32 12, i32 undef, i32 undef, i32 15>
   ret <16 x i16> %4
 }
 
 define <16 x i16> @elts_packusdw_256(<8 x i32> %a0, <8 x i32> %a1) {
 ; CHECK-LABEL: @elts_packusdw_256(
 ; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <8 x i32> %a1, <8 x i32> undef, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[TMP2:%.*]] = call <16 x i16> @llvm.x86.avx2.packusdw(<8 x i32> %a0, <8 x i32> [[TMP1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call <16 x i16> @llvm.x86.avx2.packusdw(<8 x i32> undef, <8 x i32> [[TMP1]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <16 x i16> [[TMP2]], <16 x i16> undef, <16 x i32> <i32 4, i32 5, i32 6, i32 7, i32 undef, i32 undef, i32 undef, i32 undef, i32 12, i32 13, i32 14, i32 15, i32 undef, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    ret <16 x i16> [[TMP3]]
 ;
   %1 = shufflevector <8 x i32> %a0, <8 x i32> undef, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   %2 = shufflevector <8 x i32> %a1, <8 x i32> undef, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
   %3 = call <16 x i16> @llvm.x86.avx2.packusdw(<8 x i32> %1, <8 x i32> %2)
   %4 = shufflevector <16 x i16> %3, <16 x i16> undef, <16 x i32> <i32 4, i32 5, i32 6, i32 7, i32 undef, i32 undef, i32 undef, i32 undef, i32 12, i32 13, i32 14, i32 15, i32 undef, i32 undef, i32 undef, i32 undef>
   ret <16 x i16> %4
 }
 
 define <32 x i8> @elts_packsswb_256(<16 x i16> %a0, <16 x i16> %a1) {
 ; CHECK-LABEL: @elts_packsswb_256(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <16 x i16> %a0, i16 0, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <16 x i16> %a1, i16 0, i32 8
-; CHECK-NEXT:    [[TMP3:%.*]] = call <32 x i8> @llvm.x86.avx2.packsswb(<16 x i16> [[TMP1]], <16 x i16> [[TMP2]])
-; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <32 x i8> [[TMP3]], <32 x i8> undef, <32 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24>
-; CHECK-NEXT:    ret <32 x i8> [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <32 x i8> @llvm.x86.avx2.packsswb(<16 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>, <16 x i16> <i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>)
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <32 x i8> [[TMP1]], <32 x i8> undef, <32 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24>
+; CHECK-NEXT:    ret <32 x i8> [[TMP2]]
 ;
   %1 = insertelement <16 x i16> %a0, i16 0, i32 0
   %2 = insertelement <16 x i16> %a1, i16 0, i32 8
   %3 = call <32 x i8> @llvm.x86.avx2.packsswb(<16 x i16> %1, <16 x i16> %2)
   %4 = shufflevector <32 x i8> %3, <32 x i8> undef, <32 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24, i32 24>
   ret <32 x i8> %4
 }
 
 define <32 x i8> @elts_packuswb_256(<16 x i16> %a0, <16 x i16> %a1) {
 ; CHECK-LABEL: @elts_packuswb_256(
-; CHECK-NEXT:    [[TMP1:%.*]] = call <32 x i8> @llvm.x86.avx2.packuswb(<16 x i16> <i16 undef, i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>, <16 x i16> <i16 0, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>)
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <32 x i8> [[TMP1]], <32 x i8> undef, <32 x i32> zeroinitializer
-; CHECK-NEXT:    ret <32 x i8> [[TMP2]]
+; CHECK-NEXT:    ret <32 x i8> undef
 ;
   %1 = insertelement <16 x i16> undef, i16 0, i32 1
   %2 = insertelement <16 x i16> undef, i16 0, i32 0
   %3 = call <32 x i8> @llvm.x86.avx2.packuswb(<16 x i16> %1, <16 x i16> %2)
   %4 = shufflevector <32 x i8> %3, <32 x i8> undef, <32 x i32> zeroinitializer
   ret <32 x i8> %4
 }