[InstCombine] Canonicalize insert sequences that broadcast a value into a insert + shuffle

[mkuper]

Summary

A combine that canonicalizes a chain of inserts that broadcasts a value into a single insert + a splat shufflevector.

This fixes PR31286, since we actually do the right thing when the broadcast is canonical.
(See "good4"test)

[mkuper]

A gentle holiday ping.

[majnemer]

LGTM

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

641–645 ↗

(On Diff #82119)

It might be more efficient to do ConstantAggregateZero::get(VectorType::get(...) instead.

llvm/trunk/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

for (int i = 0; i != MaskSize; ++i) {
     int Elt = Shuf.getMaskValue(i);
     if (Elt != -1 && Elt != i && Elt != i + VecSize)
       return false;
   }
 
   return true;
 }
 
+// Turn a chain of inserts that splats a value into a canonical insert + shuffle
+// splat. That is:
+// insertelt(insertelt(insertelt(insertelt X, %k, 0), %k, 1), %k, 2) ... ->
+// shufflevector(insertelt(X, %k, 0), undef, zero)
+static Instruction *foldInsSequenceIntoBroadcast(InsertElementInst &InsElt) {
+  // We are interested in the last insert in a chain. So, if this insert
+  // has a single user, and that user is an insert, bail.
+  if (InsElt.hasOneUse() && isa<InsertElementInst>(InsElt.user_back()))
+    return nullptr;
+
+  VectorType *VT = cast<VectorType>(InsElt.getType());
+  int NumElements = VT->getNumElements();
+
+  // Do not try to do this for a one-element vector, since that's a nop,
+  // and will cause an inf-loop.
+  if (NumElements == 1)
+    return nullptr;
+
+  Value *SplatVal = InsElt.getOperand(1);
+  InsertElementInst *CurrIE = &InsElt;  
+  SmallVector<bool, 16> ElementPresent(NumElements, false);
+
+  // Walk the chain backwards, keeping track of which indices we inserted into,
+  // until we hit something that isn't an insert of the splatted value.
+  while (CurrIE) {
+    ConstantInt *Idx = dyn_cast<ConstantInt>(CurrIE->getOperand(2));
+    if (!Idx || CurrIE->getOperand(1) != SplatVal)
+      return nullptr;
+
+    // Check none of the intermediate steps have any additional uses.
+    if ((CurrIE != &InsElt) && !CurrIE->hasOneUse())
+      return nullptr;
+
+    ElementPresent[Idx->getZExtValue()] = true;
+    CurrIE = dyn_cast<InsertElementInst>(CurrIE->getOperand(0));
+  }
+
+  // Make sure we've seen an insert into every element.
+  if (llvm::any_of(ElementPresent, [](bool Present) { return !Present; }))
+    return nullptr;
+
+  // All right, create the insert + shuffle.
+  Instruction *InsertFirst = InsertElementInst::Create(
+      UndefValue::get(VT), SplatVal,
+      ConstantInt::get(Type::getInt32Ty(InsElt.getContext()), 0), "", &InsElt);
+
+  Constant *ZeroMask = ConstantAggregateZero::get(
+      VectorType::get(Type::getInt32Ty(InsElt.getContext()), NumElements));
+
+  return new ShuffleVectorInst(InsertFirst, UndefValue::get(VT), ZeroMask);
+}
+
 /// insertelt (shufflevector X, CVec, Mask|insertelt X, C1, CIndex1), C, CIndex
 /// --> shufflevector X, CVec', Mask'
 static Instruction *foldConstantInsEltIntoShuffle(InsertElementInst &InsElt) {
   auto *Inst = dyn_cast<Instruction>(InsElt.getOperand(0));
   // Bail out if the parent has more than one use. In that case, we'd be
   // replacing the insertelt with a shuffle, and that's not a clear win.
   if (!Inst || !Inst->hasOneUse())
     return nullptr;
if (Value *V = SimplifyDemandedVectorElts(&IE, AllOnesEltMask, UndefElts)) {
     if (V != &IE)
       return replaceInstUsesWith(IE, V);
     return &IE;
   }
 
   if (Instruction *Shuf = foldConstantInsEltIntoShuffle(IE))
     return Shuf;
 
+  // Turn a sequence of inserts that broadcasts a scalar into a single
+  // insert + shufflevector.
+  if (Instruction *Broadcast = foldInsSequenceIntoBroadcast(IE))
+    return Broadcast;
+
   return nullptr;
 }
 
 /// Return true if we can evaluate the specified expression tree if the vector
 /// elements were shuffled in a different order.
 static bool CanEvaluateShuffled(Value *V, ArrayRef<int> Mask,
                                 unsigned Depth = 5) {
   // We can always reorder the elements of a constant.

llvm/trunk/test/Transforms/BBVectorize/X86/loop1.ll

for.body:                                         ; preds = %for.body, %entry
   %indvars.iv.next = add i64 %indvars.iv, 1
   %lftr.wideiv = trunc i64 %indvars.iv.next to i32
   %exitcond = icmp eq i32 %lftr.wideiv, 10
   br i1 %exitcond, label %for.end, label %for.body
 ; CHECK: insertelement
 ; CHECK-NEXT: insertelement
 ; CHECK-NEXT: fadd <2 x double>
 ; CHECK-NEXT: insertelement
-; CHECK-NEXT: insertelement
+; CHECK-NEXT: shufflevector
 ; CHECK-NEXT: fadd <2 x double>
 ; CHECK-NEXT: insertelement
 ; CHECK-NEXT: fmul <2 x double>
 
 ; CHECK-UNRL: %mul = fmul <2 x double> %2, %2
 ; CHECK-UNRL: %mul3 = fmul <2 x double> %2, %3
 ; CHECK-UNRL: %add = fadd <2 x double> %mul, %mul3
 ; CHECK-UNRL: %add4 = fadd <2 x double> %3, %3

llvm/trunk/test/Transforms/BBVectorize/loop1.ll

 ; CHECK: %1 = load double, double* %arrayidx2, align 8
 ; CHECK: %mul = fmul double %0, %0
 ; CHECK: %mul3 = fmul double %0, %1
 ; CHECK: %add = fadd double %mul, %mul3
 ; CHECK: %mul8 = fmul double %1, %1
 ; CHECK: %add4.v.i1.1 = insertelement <2 x double> undef, double %1, i32 0
 ; CHECK: %add4.v.i1.2 = insertelement <2 x double> %add4.v.i1.1, double %0, i32 1
 ; CHECK: %add4 = fadd <2 x double> %add4.v.i1.2, %add4.v.i1.2
-; CHECK: %add5.v.i1.1 = insertelement <2 x double> undef, double %0, i32 0
-; CHECK: %add5.v.i1.2 = insertelement <2 x double> %add5.v.i1.1, double %0, i32 1
+; CHECK: %2 = insertelement <2 x double> undef, double %0, i32 0
+; CHECK: %add5.v.i1.2 = shufflevector <2 x double> %2, <2 x double> undef, <2 x i32> zeroinitializer
 ; CHECK: %add5 = fadd <2 x double> %add4, %add5.v.i1.2
-; CHECK: %mul6.v.i0.2 = insertelement <2 x double> %add5.v.i1.1, double %mul8, i32 1
+; CHECK: %mul6.v.i0.2 = insertelement <2 x double> %2, double %mul8, i32 1
 ; CHECK: %mul6 = fmul <2 x double> %mul6.v.i0.2, %add5
 ; CHECK: %mul6.v.r1 = extractelement <2 x double> %mul6, i32 0
 ; CHECK: %mul6.v.r2 = extractelement <2 x double> %mul6, i32 1
 ; CHECK: %add7 = fadd double %add, %mul6.v.r1
 ; CHECK: %add12 = fadd double %add7, %mul6.v.r2
 ; CHECK: %arrayidx14 = getelementptr inbounds double, double* %out, i64 %indvars.iv
 ; CHECK: store double %add12, double* %arrayidx14, align 8
 ; CHECK: %indvars.iv.next = add i64 %indvars.iv, 1

llvm/trunk/test/Transforms/BBVectorize/simple-int.ll

define double @testcopysign(double %A1, double %A2, double %B1, double %B2) {
   %R  = fmul double %Z1, %Z2
   ret double %R
 ; CHECK: @testcopysign
 ; CHECK: %X1.v.i1.1 = insertelement <2 x double> undef, double %B1, i32 0
 ; CHECK: %X1.v.i1.2 = insertelement <2 x double> %X1.v.i1.1, double %B2, i32 1
 ; CHECK: %X1.v.i0.1 = insertelement <2 x double> undef, double %A1, i32 0
 ; CHECK: %X1.v.i0.2 = insertelement <2 x double> %X1.v.i0.1, double %A2, i32 1
 ; CHECK: %X1 = fsub <2 x double> %X1.v.i0.2, %X1.v.i1.2
-; CHECK: %Y1.v.i1.2 = insertelement <2 x double> %X1.v.i0.1, double %A1, i32 1
+; CHECK: %Y1.v.i1.2 = shufflevector <2 x double> %X1.v.i0.1, <2 x double> undef, <2 x i32> zeroinitializer
 ; CHECK: %Y1 = call <2 x double> @llvm.copysign.v2f64(<2 x double> %X1, <2 x double> %Y1.v.i1.2)
 ; CHECK: %Z1 = fadd <2 x double> %Y1, %X1.v.i1.2
 ; CHECK: %Z1.v.r1 = extractelement <2 x double> %Z1, i32 0
 ; CHECK: %Z1.v.r2 = extractelement <2 x double> %Z1, i32 1
 ; CHECK: %R = fmul double %Z1.v.r1, %Z1.v.r2
 ; CHECK: ret double %R
 
 }

llvm/trunk/test/Transforms/InstCombine/broadcast.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -instcombine -S < %s | FileCheck %s
+
+; CHECK-LABEL: good1
+; CHECK: %[[INS:.*]] = insertelement <4 x float> undef, float %arg, i32 0
+; CHECK-NEXT: %[[BCAST:.*]] = shufflevector <4 x float> %[[INS]], <4 x float> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT: ret <4 x float> %[[BCAST]]
+define <4 x float> @good1(float %arg) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: good2
+; CHECK: %[[INS:.*]] = insertelement <4 x float> undef, float %arg, i32 0
+; CHECK-NEXT: %[[BCAST:.*]] = shufflevector <4 x float> %[[INS]], <4 x float> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT: ret <4 x float> %[[BCAST]]
+define <4 x float> @good2(float %arg) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 1
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 2
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 0
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: good3
+; CHECK: %[[INS:.*]] = insertelement <4 x float> undef, float %arg, i32 0
+; CHECK-NEXT: %[[BCAST:.*]] = shufflevector <4 x float> %[[INS]], <4 x float> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT: ret <4 x float> %[[BCAST]]
+define <4 x float> @good3(float %arg) {
+  %tmp = insertelement <4 x float> zeroinitializer, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: good4
+; CHECK: %[[INS:.*]] = insertelement <4 x float> undef, float %arg, i32 0
+; CHECK-NEXT: %[[ADD:.*]] = fadd <4 x float> %[[INS]], %[[INS]]
+; CHECK-NEXT: %[[BCAST:.*]] = shufflevector <4 x float> %[[ADD]], <4 x float> undef, <4 x i32> zeroinitializer
+; CHECK-NEXT: ret <4 x float> %[[BCAST]]
+define <4 x float> @good4(float %arg) {
+  %tmp = insertelement <4 x float> zeroinitializer, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  %tmp7 = fadd <4 x float> %tmp6, %tmp6
+  ret <4 x float> %tmp7
+}
+
+; CHECK-LABEL: bad1
+; CHECK-NOT: shufflevector
+define <4 x float> @bad1(float %arg) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 1
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: bad2
+; CHECK-NOT: shufflevector
+define <4 x float> @bad2(float %arg) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 0
+  %tmp5 = insertelement <4 x float> %tmp, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: bad3
+; CHECK-NOT: shufflevector
+define <4 x float> @bad3(float %arg, float %arg2) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg2, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}
+
+; CHECK-LABEL: bad4
+; CHECK-NOT: shufflevector
+define <1 x float> @bad4(float %arg) {
+  %tmp = insertelement <1 x float> undef, float %arg, i32 0
+  ret <1 x float> %tmp
+}
+
+; CHECK-LABEL: bad5
+; CHECK-NOT: shufflevector
+define <4 x float> @bad5(float %arg) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 2
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  %tmp7 = fadd <4 x float> %tmp6, %tmp4
+  ret <4 x float> %tmp7
+}
+
+; CHECK-LABEL: bad6
+; CHECK-NOT: shufflevector
+define <4 x float> @bad6(float %arg, i32 %k) {
+  %tmp = insertelement <4 x float> undef, float %arg, i32 0
+  %tmp4 = insertelement <4 x float> %tmp, float %arg, i32 1
+  %tmp5 = insertelement <4 x float> %tmp4, float %arg, i32 %k
+  %tmp6 = insertelement <4 x float> %tmp5, float %arg, i32 3
+  ret <4 x float> %tmp6
+}

llvm/trunk/test/Transforms/SLPVectorizer/X86/operandorder.ll

define void @shuffle_operands1(double * noalias %from, double * noalias %to,
   %to_2 = getelementptr double, double * %to, i64 1
   store double %v1_1, double *%to
   store double %v1_2, double *%to_2
   ret void
 }
 
 ; CHECK-LABEL: shuffle_preserve_broadcast
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]
   %from_1 = getelementptr double, double *%from, i64 1
   %v0_1 = load double , double * %from
   %v0_2 = load double , double * %from_1
   %v1_1 = fadd double %v0_1, %p
   %v1_2 = fadd double %v0_1, %v0_2
   %to_2 = getelementptr double, double * %to, i64 1
   store double %v1_1, double *%to
   store double %v1_2, double *%to_2
 br i1 undef, label %lp, label %ext
 
 ext:
   ret void
 }
 
 ; CHECK-LABEL: shuffle_preserve_broadcast2
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast2(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]
   %from_1 = getelementptr double, double *%from, i64 1
   %v0_1 = load double , double * %from
   %v0_2 = load double , double * %from_1
   %v1_1 = fadd double %p, %v0_1
   %v1_2 = fadd double %v0_2, %v0_1
   %to_2 = getelementptr double, double * %to, i64 1
   store double %v1_1, double *%to
   store double %v1_2, double *%to_2
 br i1 undef, label %lp, label %ext
 
 ext:
   ret void
 }
 
 ; CHECK-LABEL: shuffle_preserve_broadcast3
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast3(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]

 
 ext:
   ret void
 }
 
 
 ; CHECK-LABEL: shuffle_preserve_broadcast4
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast4(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]
   %from_1 = getelementptr double, double *%from, i64 1
   %v0_1 = load double , double * %from
   %v0_2 = load double , double * %from_1
   %v1_1 = fadd double %v0_2, %v0_1
   %v1_2 = fadd double %p, %v0_1
   %to_2 = getelementptr double, double * %to, i64 1
   store double %v1_1, double *%to
   store double %v1_2, double *%to_2
 br i1 undef, label %lp, label %ext
 
 ext:
   ret void
 }
 
 ; CHECK-LABEL: shuffle_preserve_broadcast5
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast5(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]

 
 ext:
   ret void
 }
 
 
 ; CHECK-LABEL: shuffle_preserve_broadcast6
 ; CHECK: %[[BCAST:[a-z0-9]+]] = insertelement <2 x double> undef, double %v0_1
-; CHECK:                      = insertelement <2 x double> %[[BCAST]], double %v0_1
+; CHECK:                      = shufflevector <2 x double> %[[BCAST]], <2 x double> undef, <2 x i32> zeroinitializer
 define void @shuffle_preserve_broadcast6(double * noalias %from,
                                         double * noalias %to,
                                         double %v1, double %v2) {
 entry:
 br label %lp
 
 lp:
   %p = phi double [ 1.000000e+00, %lp ], [ 0.000000e+00, %entry ]