[InstCombine] shrink truncated splat shuffle

[spatel]

Summary

This could be one part of solving a recent bug report:
http://lists.llvm.org/pipermail/llvm-dev/2017-February/110293.html

This keeps with our general approach: changing arbitrary shuffles is off-limts, but changing splat is ok. The transform is very similar to the existing shrinkBitwiseLogic() canonicalization.

[efriedma]

"the mask doesn't change"

That's technically true... but doesn't really match the spirit of the rule that we don't change shuffle masks. Many targets have shuffles that only work at specific widths.

[spatel]

In D30123#680573, @efriedma wrote:

"the mask doesn't change"

That's technically true... but doesn't really match the spirit of the rule that we don't change shuffle masks. Many targets have shuffles that only work at specific widths.

I knew I might be pushing the limits here by generalizing too much. :)
In the motivating example, we have a splat mask, so we could still do a limited form of this transform and catch that case. Can we assume that any splat shuffle is safe to transform in target-independent IR?

[efriedma]

instcombine already assumes that targets can handle a splat of any type; that's fine.

[spatel]

Patch upated:

1. Limit the transform to splat shuffles. In this case, we canonicalize the vector to operand 0 (and operand 1 should be undef), so that eliminates the commuted pattern match.
2. Added test for splat with strange types for more coverage of the getSplatValue() implementation.
3. Added test with undef element in mask. Unlike the DAG's ShuffleVectorSDNode::isSplatMask(), Constant::getSplatValue() does not check for undef elements, so we miss this case currently. I'm not sure if this matters in practice currently because we're conservative about changing shuffle masks.

[spatel]

Ping.

[spatel]

Ping * 2.

[majnemer]

@efriedma, @mkuper I think you two are well positioned to review this change. Please take a look when you get a chance. Thanks!

[efriedma]

LGTM, but you might want to look into the following...

define <8 x i16> @wide_splat1(<8 x i32> %x) {
  %shuf = trunc <8 x i32> %x to <8 x i16>
  %trunc = shufflevector <8 x i16> %shuf, <8 x i16> undef, <8 x i32> <i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2>
  ret <8 x i16> %trunc
}

define <8 x i16> @wide_splat2(<8 x i32> %x) {
  %trunc = shufflevector <8 x i32> %x, <8 x i32> undef, <8 x i32> <i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2>
  %shuf = trunc <8 x i32> %trunc to <8 x i16>
  ret <8 x i16> %shuf
}

llc output (x86-64, no SSSE3):

        .text
        .file   "<stdin>"
        .globl  wide_splat1
        .p2align        4, 0x90
        .type   wide_splat1,@function
wide_splat1:                            # @wide_splat1
        .cfi_startproc
# BB#0:
        pslld   $16, %xmm1
        psrad   $16, %xmm1
        pslld   $16, %xmm0
        psrad   $16, %xmm0
        packssdw        %xmm1, %xmm0
        pshuflw $170, %xmm0, %xmm0      # xmm0 = xmm0[2,2,2,2,4,5,6,7]
        pshufd  $80, %xmm0, %xmm0       # xmm0 = xmm0[0,0,1,1]
        retq
.Lfunc_end0:
        .size   wide_splat1, .Lfunc_end0-wide_splat1
        .cfi_endproc

        .globl  wide_splat2
        .p2align        4, 0x90
        .type   wide_splat2,@function
wide_splat2:                            # @wide_splat2
        .cfi_startproc
# BB#0:
        pshufd  $170, %xmm0, %xmm0      # xmm0 = xmm0[2,2,2,2]
        pslld   $16, %xmm0
        psrad   $16, %xmm0
        packssdw        %xmm0, %xmm0
        retq
.Lfunc_end1:
        .size   wide_splat2, .Lfunc_end1-wide_splat2
        .cfi_endproc

[spatel]

In D30123#693699, @efriedma wrote:

LGTM, but you might want to look into the following...

define <8 x i16> @wide_splat1(<8 x i32> %x) {
  %shuf = trunc <8 x i32> %x to <8 x i16>
  %trunc = shufflevector <8 x i16> %shuf, <8 x i16> undef, <8 x i32> <i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2>
  ret <8 x i16> %trunc
}

Thanks for pointing it out. Filed as:
https://bugs.llvm.org/show_bug.cgi?id=32160

lib/Transforms/InstCombine/InstCombineCasts.cpp

if (!match(Trunc.getOperand(0), m_OneUse(m_BinOp(LogicOp))) ||
     return nullptr;
 
   // trunc (logic X, C) --> logic (trunc X, C')
   Constant *NarrowC = ConstantExpr::getTrunc(C, DestTy);
   Value *NarrowOp0 = Builder->CreateTrunc(LogicOp->getOperand(0), DestTy);
   return BinaryOperator::Create(LogicOp->getOpcode(), NarrowOp0, NarrowC);
 }
 
+/// Try to narrow the width of a splat shuffle. This could be generalized to any
+/// shuffle with a constant operand, but we limit the transform to avoid
+/// creating a shuffle type that targets may not be able to lower effectively.
+static Instruction *shrinkSplatShuffle(TruncInst &Trunc,
+                                       InstCombiner::BuilderTy &Builder) {
+  auto *Shuf = dyn_cast<ShuffleVectorInst>(Trunc.getOperand(0));
+  if (Shuf && Shuf->hasOneUse() && isa<UndefValue>(Shuf->getOperand(1)) &&
+      Shuf->getMask()->getSplatValue()) {
+    // trunc (shuf X, Undef, SplatMask) --> shuf (trunc X), Undef, SplatMask
+    Constant *NarrowUndef = UndefValue::get(Trunc.getType());
+    Value *NarrowOp = Builder.CreateTrunc(Shuf->getOperand(0), Trunc.getType());
+    return new ShuffleVectorInst(NarrowOp, NarrowUndef, Shuf->getMask());
+  }
+
+  return nullptr;
+}
+
 Instruction *InstCombiner::visitTrunc(TruncInst &CI) {
   if (Instruction *Result = commonCastTransforms(CI))
     return Result;
 
   // Test if the trunc is the user of a select which is part of a
   // minimum or maximum operation. If so, don't do any more simplification.
   // Even simplifying demanded bits can break the canonical form of a
   // min/max.
if (Cst->getValue().ult(ASize)) {
       Shift->takeName(Src);
       return CastInst::CreateIntegerCast(Shift, CI.getType(), true);
     }
   }
 
   if (Instruction *I = shrinkBitwiseLogic(CI))
     return I;
 
+  if (Instruction *I = shrinkSplatShuffle(CI, *Builder))
+    return I;
+
   if (Src->hasOneUse() && isa<IntegerType>(SrcTy) &&
       shouldChangeType(SrcTy, DestTy)) {
     // Transform "trunc (shl X, cst)" -> "shl (trunc X), cst" so long as the
     // dest type is native and cst < dest size.
     if (match(Src, m_Shl(m_Value(A), m_ConstantInt(Cst))) &&
         !match(A, m_Shr(m_Value(), m_Constant()))) {
       // Skip shifts of shift by constants. It undoes a combine in
       // FoldShiftByConstant and is the extend in reg pattern.

test/Transforms/InstCombine/trunc.ll

 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i8>
 ; CHECK-NEXT:    ret <4 x i8> [[TRUNC]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> <i32 35, i32 3634, i32 90, i32 -1>, <4 x i32> <i32 1, i32 5, i32 6, i32 2>
   %trunc = trunc <4 x i32> %shuf to <4 x i8>
   ret <4 x i8> %trunc
 }
 
-; FIXME:
 ; trunc (shuffle X, undef, SplatMask) --> shuffle (trunc X), undef, SplatMask
 
 define <4 x i8> @wide_splat1(<4 x i32> %x) {
 ; CHECK-LABEL: @wide_splat1(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i8>
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc <4 x i32> %x to <4 x i8>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shufflevector <4 x i8> [[TMP1]], <4 x i8> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
 ; CHECK-NEXT:    ret <4 x i8> [[TRUNC]]
 ;
   %shuf = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
   %trunc = trunc <4 x i32> %shuf to <4 x i8>
   ret <4 x i8> %trunc
 }
 
-; FIXME:
 ; Test weird types.
 ; trunc (shuffle X, undef, SplatMask) --> shuffle (trunc X), undef, SplatMask
 
 define <3 x i31> @wide_splat2(<3 x i33> %x) {
 ; CHECK-LABEL: @wide_splat2(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <3 x i33> %x, <3 x i33> undef, <3 x i32> <i32 1, i32 1, i32 1>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <3 x i33> [[SHUF]] to <3 x i31>
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc <3 x i33> %x to <3 x i31>
+; CHECK-NEXT:    [[TRUNC:%.*]] = shufflevector <3 x i31> [[TMP1]], <3 x i31> undef, <3 x i32> <i32 1, i32 1, i32 1>
 ; CHECK-NEXT:    ret <3 x i31> [[TRUNC]]
 ;
   %shuf = shufflevector <3 x i33> %x, <3 x i33> undef, <3 x i32> <i32 1, i32 1, i32 1>
   %trunc = trunc <3 x i33> %shuf to <3 x i31>
   ret <3 x i31> %trunc
 }
 
 ; FIXME: