[InstCombine] Fix for PR29124: reduce insertelements to shufflevector

[ABataev]

Summary

If inserting more than one constant into a vector:

define <4 x float> @foo(<4 x float> %x) {
  %ins1 = insertelement <4 x float> %x, float 1.0, i32 1
  %ins2 = insertelement <4 x float> %ins1, float 2.0, i32 2
  ret <4 x float> %ins2
}

InstCombine could reduce that to a shufflevector:

define <4 x float> @goo(<4 x float> %x) {
 %shuf = shufflevector <4 x float> %x, <4 x float> <float undef, float 1.0, float 2.0, float undef>, <4 x i32><i32 0, i32 5, i32 6, i32 3>
 ret <4 x float> %shuf
}

[spatel]

Adding reviewers/subscribers from D23886.

If D23886 goes in, I think this patch can be simplified to just work on a pair of insertelements with an undef vector operand for the first insert. Ie, this transform provides the 'seed' shufflevector, and the transform in D23886 pulls subsequent inserts into that shuffle.

test/Transforms/InstCombine/insert-const-shuf.ll

17–22 ↗

(On Diff #70142)

This will be affected by D23886. Ideally, this should be only one shuffle, so this means we are missing a fold that combines shufflevectors.

[spatel]

In D24182#533000, @spatel wrote:

If D23886 goes in, I think this patch can be simplified to just work on a pair of insertelements with an undef vector operand for the first insert. Ie, this transform provides the 'seed' shufflevector, and the transform in D23886 pulls subsequent inserts into that shuffle.

Hi Alexey -

Does your implementation catch cases that the above idea would not? Note that D23866 was committed at rL280504, and I updated the checks in vec_demanded_elts.ll at rL280531 because it was hard to tell what was changing. So you will need to merge this patch after those changes.

Also, I filed https://llvm.org/bugs/show_bug.cgi?id=30264 for the missing shuffle/shuffle fold (although it may not be important if we have this combine + D23866).

Finally, note Hal's comment ( https://llvm.org/bugs/show_bug.cgi?id=30264#c1 ) about not creating risky shuffles; this is something that Eli mentioned in D23866 as well. I think the case of combining two insertelt constants would be safe for any target (AFAICT, the legalization fallback is to just create inserts again), but someone can tell me if I'm wrong about that. :)

[RKSimon]

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

670

Is this actually firing? Might it be better to just bail out if this happens?

test/Transforms/InstCombine/insert-const-shuf.ll

13 ↗

(On Diff #70142)

Shouldn't D23886 have dealt with this?

test/Transforms/InstCombine/vec_demanded_elts.ll

1

Possibly regenerate this test first so that the diff is more obvious?

[ABataev]

Updated to the latest version + small fixes

[ABataev]

Full diff

[ABataev]

Fixed regression in the optimization

[ABataev]

Updated to latest version, renamed UndefElts3/4 to (L|R)HSUndefElts, joined with D24381 to avoid regressions

[ABataev]

Added a test

[ABataev]

Added tests with constants + small improvements

[RKSimon]

LGTM with 2 minor corrections.

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

658

Please add a comment describing what this loop is doing.

667

Please add a comment describing what this loop is doing.

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

   return true;
 }
 
-/// insertelt (shufflevector X, CVec, Mask), C, CIndex -->
-/// shufflevector X, CVec', Mask'
+/// insertelt (shufflevector X, CVec, Mask|insertelt X, C1, CIndex1), C, CIndex
+/// --> shufflevector X, CVec', Mask'
 static Instruction *foldConstantInsEltIntoShuffle(InsertElementInst &InsElt) {
-  // Bail out if the shuffle has more than one use. In that case, we'd be
+  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.
-  auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0));
-  if (!Shuf || !Shuf->hasOneUse())
+  if (!Inst || !Inst->hasOneUse())
     return nullptr;
+  Value *X = nullptr;
+  Constant *NewShufVec = nullptr;
+  Constant *NewMask = nullptr;
+  if (auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0))) {
+    // The shuffle must have a constant vector operand. The insertelt must have
+    // a constant scalar being inserted at a constant position in the vector.
+    Constant *ShufConstVec, *InsEltScalar;
+    uint64_t InsEltIndex;
+    if (!match(Shuf->getOperand(1), m_Constant(ShufConstVec)) ||
+        !match(InsElt.getOperand(1), m_Constant(InsEltScalar)) ||
+        !match(InsElt.getOperand(2), m_ConstantInt(InsEltIndex)))
+      return nullptr;
 
-  // The shuffle must have a constant vector operand. The insertelt must have a
-  // constant scalar being inserted at a constant position in the vector.
-  Constant *ShufConstVec, *InsEltScalar;
-  uint64_t InsEltIndex;
-  if (!match(Shuf->getOperand(1), m_Constant(ShufConstVec)) ||
-      !match(InsElt.getOperand(1), m_Constant(InsEltScalar)) ||
-      !match(InsElt.getOperand(2), m_ConstantInt(InsEltIndex)))
-    return nullptr;
+    // Adding an element to an arbitrary shuffle could be expensive, but a
+    // shuffle that selects elements from vectors without crossing lanes is
+    // assumed cheap.
+    // If we're just adding a constant into that shuffle, it will still be
+    // cheap.
+    if (!isShuffleEquivalentToSelect(*Shuf))
+      return nullptr;
 
-  // Adding an element to an arbitrary shuffle could be expensive, but a shuffle
-  // that selects elements from vectors without crossing lanes is assumed cheap.
-  // If we're just adding a constant into that shuffle, it will still be cheap.
-  if (!isShuffleEquivalentToSelect(*Shuf))
-    return nullptr;
+    // From the above 'select' check, we know that the mask has the same number
+    // of elements as the vector input operands. We also know that each constant
+    // input element is used in its lane and can not be used more than once by
+    // the shuffle. Therefore, replace the constant in the shuffle's constant
+    // vector with the insertelt constant. Replace the constant in the shuffle's
+    // mask vector with the insertelt index plus the length of the vector
+    // (because the constant vector operand of a shuffle is always the 2nd
+    // operand).
+    Constant *Mask = Shuf->getMask();
+    unsigned NumElts = Mask->getType()->getVectorNumElements();
+    SmallVector<Constant *, 16> NewShufElts(NumElts);
+    SmallVector<Constant *, 16> NewMaskElts(NumElts);
+    for (unsigned I = 0; I != NumElts; ++I) {
+      if (I == InsEltIndex) {
+        NewShufElts[I] = InsEltScalar;
+        Type *Int32Ty = Type::getInt32Ty(Shuf->getContext());
+        NewMaskElts[I] = ConstantInt::get(Int32Ty, InsEltIndex + NumElts);
+      } else {
+        // Copy over the existing values.
+        NewShufElts[I] = ShufConstVec->getAggregateElement(I);
+        NewMaskElts[I] = Mask->getAggregateElement(I);
+      }
+    }
 
-  // From the above 'select' check, we know that the mask has the same number of
-  // elements as the vector input operands. We also know that each constant
-  // input element is used in its lane and can not be used more than once by the
-  // shuffle. Therefore, replace the constant in the shuffle's constant vector
-  // with the insertelt constant. Replace the constant in the shuffle's mask
-  // vector with the insertelt index plus the length of the vector (because the
-  // constant vector operand of a shuffle is always the 2nd operand).
-  Constant *Mask = Shuf->getMask();
-  unsigned NumElts = Mask->getType()->getVectorNumElements();
-  SmallVector<Constant*, 16> NewShufElts(NumElts);
-  SmallVector<Constant*, 16> NewMaskElts(NumElts);
-  for (unsigned i = 0; i != NumElts; ++i) {
-    if (i == InsEltIndex) {
-      NewShufElts[i] = InsEltScalar;
-      Type *Int32Ty = Type::getInt32Ty(Shuf->getContext());
-      NewMaskElts[i] = ConstantInt::get(Int32Ty, InsEltIndex + NumElts);
-    } else {
-      // Copy over the existing values.
-      NewShufElts[i] = ShufConstVec->getAggregateElement(i);
-      NewMaskElts[i] = Mask->getAggregateElement(i);
+    // Create new operands for a shuffle that includes the constant of the
+    // original insertelt. The old shuffle will be dead now.
+    X = Shuf->getOperand(0);
+    NewShufVec = ConstantVector::get(NewShufElts);
+    NewMask = ConstantVector::get(NewMaskElts);
+  } else if (auto *IEI = dyn_cast<InsertElementInst>(Inst)) {
+    // Transform sequences of insertelements ops with constant data/indexes into
+    // a single shuffle op.
+    unsigned NumElts = InsElt.getType()->getNumElements();
+
+    uint64_t InsertIdx[2];
+    Constant *Val[2];
+    if (!match(InsElt.getOperand(2), m_ConstantInt(InsertIdx[0])) ||
+        !match(InsElt.getOperand(1), m_Constant(Val[0])) ||
+        !match(IEI->getOperand(2), m_ConstantInt(InsertIdx[1])) ||
+        !match(IEI->getOperand(1), m_Constant(Val[1])))
+      return nullptr;
+    SmallVector<Constant *, 16> Values(NumElts);
+    SmallVector<Constant *, 16> Mask(NumElts);
+    auto ValI = std::begin(Val);
+    for (uint64_t I : InsertIdx) {
+      if (!Values[I]) {
+        assert(!Mask[I]);
+        Values[I] = *ValI;
+        Mask[I] = ConstantInt::get(Type::getInt32Ty(InsElt.getContext()),
+                                   NumElts + I);
+      }
+      ++ValI;
     }
+    for (unsigned I = 0; I < NumElts; ++I) {
+      if (!Values[I]) {
+        assert(!Mask[I]);
+        Values[I] = UndefValue::get(InsElt.getType()->getElementType());
+        Mask[I] = ConstantInt::get(Type::getInt32Ty(InsElt.getContext()), I);
+      }
+    }
+    // Create new operands for a shuffle that includes the constant of the
+    // original insertelt.
+    X = IEI->getOperand(0);
+    NewShufVec = ConstantVector::get(Values);
+    NewMask = ConstantVector::get(Mask);
   }
-
-  // Create new operands for a shuffle that includes the constant of the
-  // original insertelt. The old shuffle will be dead now.
-  Constant *NewShufVec = ConstantVector::get(NewShufElts);
-  Constant *NewMask = ConstantVector::get(NewMaskElts);
-  return new ShuffleVectorInst(Shuf->getOperand(0), NewShufVec, NewMask);
+  if (X && NewShufVec && NewMask)
+    return new ShuffleVectorInst(X, NewShufVec, NewMask);
+  return nullptr;
 }
 
 Instruction *InstCombiner::visitInsertElementInst(InsertElementInst &IE) {

test/Transforms/InstCombine/vec_demanded_elts.ll

 
 define <2 x float> @test_fptrunc(double %f) {
 ; CHECK-LABEL: @test_fptrunc(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x double> undef, double %f, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x double> [[TMP1]], double 0.000000e+00, i32 1
-; CHECK-NEXT:    [[TMP3:%.*]] = fptrunc <2 x double> [[TMP2]] to <2 x float>
-; CHECK-NEXT:    ret <2 x float> [[TMP3]]
+; CHECK-NEXT:    [[TMP9:%.*]] = insertelement <4 x double> undef, double %f, i32 0
+; CHECK-NEXT:    [[TMP12:%.*]] = shufflevector <4 x double> [[TMP9]], <4 x double> <double undef, double 0.000000e+00, double undef, double undef>, <4 x i32> <i32 0, i32 5, i32 undef, i32 undef>
+; CHECK-NEXT:    [[TMP5:%.*]] = fptrunc <4 x double> [[TMP12]] to <4 x float>
+; CHECK-NEXT:    [[RET:%.*]] = shufflevector <4 x float> [[TMP5]], <4 x float> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    ret <2 x float> [[RET]]
 ;
   %tmp9 = insertelement <4 x double> undef, double %f, i32 0
   %tmp10 = insertelement <4 x double> %tmp9, double 0.000000e+00, i32 1

 
 define <2 x double> @test_fpext(float %f) {
 ; CHECK-LABEL: @test_fpext(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x float> undef, float %f, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x float> [[TMP1]], float 0.000000e+00, i32 1
-; CHECK-NEXT:    [[TMP3:%.*]] = fpext <2 x float> [[TMP2]] to <2 x double>
-; CHECK-NEXT:    ret <2 x double> [[TMP3]]
+; CHECK-NEXT:    [[TMP9:%.*]] = insertelement <4 x float> undef, float %f, i32 0
+; CHECK-NEXT:    [[TMP12:%.*]] = shufflevector <4 x float> [[TMP9]], <4 x float> <float undef, float 0.000000e+00, float undef, float undef>, <4 x i32> <i32 0, i32 5, i32 undef, i32 undef>
+; CHECK-NEXT:    [[TMP5:%.*]] = fpext <4 x float> [[TMP12]] to <4 x double>
+; CHECK-NEXT:    [[RET:%.*]] = shufflevector <4 x double> [[TMP5]], <4 x double> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    ret <2 x double> [[RET]]
 ;
   %tmp9 = insertelement <4 x float> undef, float %f, i32 0
   %tmp10 = insertelement <4 x float> %tmp9, float 0.000000e+00, i32 1

 define <4 x float> @test_select(float %f, float %g) {
 ; CHECK-LABEL: @test_select(
 ; CHECK-NEXT:    [[A0:%.*]] = insertelement <4 x float> undef, float %f, i32 0
-; CHECK-NEXT:    [[A3:%.*]] = insertelement <4 x float> [[A0]], float 3.000000e+00, i32 3
+; CHECK-NEXT:    [[A3:%.*]] = shufflevector <4 x float> [[A0]], <4 x float> <float undef, float undef, float undef, float 3.000000e+00>, <4 x i32> <i32 0, i32 undef, i32 undef, i32 7>
 ; CHECK-NEXT:    [[RET:%.*]] = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x float> [[A3]], <4 x float> <float undef, float 4.000000e+00, float 5.000000e+00, float undef>
 ; CHECK-NEXT:    ret <4 x float> [[RET]]
 ;

test/Transforms/InstCombine/vector_insertelt_shuffle.ll

   ret<4 x float> %ins2
 }
 
-; FIXME: insertelements should fold to shuffle
+; insertelements should fold to shuffle
 ; CHECK-LABEL: @foo
-; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 1.000000e+00, i32 1
-; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 2.000000e+00, i32 2
+; CHECK-NEXT: shufflevector <4 x float> %{{.+}}, <4 x float> <float undef, float 1.000000e+00, float 2.000000e+00, float undef>, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT: ret <4 x float> %
 
 define<4 x float> @bar(<4 x float> %x, float %a) {

   ret<4 x float> %ins6
 }
 
-; FIXME: insertelements should fold to shuffle
+; insertelements should fold to shuffle
 ; CHECK-LABEL: @bazz
 ; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 1.000000e+00, i32 3
 ; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 5.000000e+00, i32 %
-; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 1.000000e+00, i32 1
-; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 2.000000e+00, i32 2
+; CHECK-NEXT: shufflevector <4 x float> %{{.+}}, <4 x float> <float undef, float 1.000000e+00, float 2.000000e+00, float undef>, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT: insertelement <4 x float> %{{.+}}, float 7.000000e+00, i32 %
 ; CHECK-NEXT: ret <4 x float> %