[InstCombine] fold insertelement of scalar constant into vector of constants

[spatel]

Summary

Keeping the various operand names straight for a pair of insertelements is challenging. I've reused the existing code in the patch, but if I renamed those code variables, it might be something like this:
insertelement (insertelement VectorC, X, Idx1C), ScalarC, Idx0C --> insertelement NewVectorC, X, Idx1C

So we're putting the scalar constant from the 2nd instruction into the vector constant operand of the 1st instruction because that eliminates the 2nd insertelement.

[efriedma]

Would it make sense for this transform to look through more than one insertelement? Alternatively, we could perform this transform unconditionally: transform (insertelement (insertelement, X, Y, IdxC1), C, IdxC2) -> (insertelement (insertelement X, C, IdxC2), Y, IdxC1), ignoring the actual value of X.

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

745

Reusing Val[1] like this is confusing; can you use a separate variable?

758

getAggregateElement can return null.

[spatel]

In D31196#706638, @efriedma wrote:

Alternatively, we could perform this transform unconditionally: transform (insertelement (insertelement, X, Y, IdxC1), C, IdxC2) -> (insertelement (insertelement X, C, IdxC2), Y, IdxC1), ignoring the actual value of X.

That looks great - we canonicalize harder, the constant folding will happen via the Builder, and it makes the patch cleaner.

[spatel]

Patch updated as suggested by Eli:
Canonicalize an insert of a constant before an insert of a variable. There was already a test for this (but did nothing before), so I added a comment to explain.

[efriedma]

LGTM.

(I guess this algorithm is theoretically quadratic over a long sequence of insertelements, but the constant factor is small enough that it shouldn't matter.)

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

if (auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0))) {
 
     // Create new operands for a shuffle that includes the constant of the
     // original insertelt. The old shuffle will be dead now.
     return new ShuffleVectorInst(Shuf->getOperand(0),
                                  ConstantVector::get(NewShufElts),
                                  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.
+    // a single shuffle op or a single insertelement.
     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])))
+    if (!match(InsElt.getOperand(1), m_Constant(Val[0])) ||
+        !match(InsElt.getOperand(2), m_ConstantInt(InsertIdx[0])) ||
+        !match(IEI->getOperand(2), m_ConstantInt(InsertIdx[1])))
       return nullptr;
+
+    if (match(IEI->getOperand(1), m_Constant(Val[1]))) {
+      // inselt1 (inselt2 X, Val1C, Idx1C), Val0C, Idx0C --> shuffle
       SmallVector<Constant *, 16> Values(NumElts);
       SmallVector<Constant *, 16> Mask(NumElts);
       auto ValI = std::begin(Val);
       // Generate new constant vector and mask.
-    // We have 2 values/masks from the insertelements instructions. Insert them
-    // into new value/mask vectors.
+      // We have 2 values/masks from the insertelements instructions. Insert
+      // them into new value/mask vectors.
       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;
       }
       // Remaining values are filled with 'undef' values.
       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.
       return new ShuffleVectorInst(IEI->getOperand(0),
                                    ConstantVector::get(Values),
                                    ConstantVector::get(Mask));
     }
+    if (match(IEI->getOperand(0), m_Constant(Val[1])) &&
+        InsertIdx[0] != InsertIdx[1]) {
+      // If we're inserting a scalar constant into a vector of constants with
+      // one variable element, eliminate the 2nd insert by putting the scalar
+      // constant directly into the vector constant:
+      // ins0 (ins1 Val1C, X, Idx1C), Val0C, Idx0C --> ins Val1C', X, Idx1C
+      SmallVector<Constant *, 16> NewVectorC(NumElts);
+      for (unsigned I = 0; I != NumElts; ++I) {
+        if (I == InsertIdx[0])
+          // Replace the old element with the constant from the 2nd insert.
+          NewVectorC[I] = Val[0];
+        else
+          // Copy over the existing values.
+          NewVectorC[I] = Val[1]->getAggregateElement(I);
+      }
+      return InsertElementInst::Create(ConstantVector::get(NewVectorC),
+                                       IEI->getOperand(1), IEI->getOperand(2));
+    }
+  }
   return nullptr;
 }
 
 Instruction *InstCombiner::visitInsertElementInst(InsertElementInst &IE) {
   Value *VecOp    = IE.getOperand(0);
   Value *ScalarOp = IE.getOperand(1);
   Value *IdxOp    = IE.getOperand(2);
 

test/Transforms/InstCombine/bitcast-bigendian.ll

;
   %tmp33 = shl i64 %tmp32, 32
   %ins35 = or i64 %tmp33, %tmp38
   %tmp43 = bitcast i64 %ins35 to <2 x float>
   ret <2 x float> %tmp43
 }
 
 define <2 x float> @test6(float %A){
 ; CHECK-LABEL: @test6(
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x float> undef, float %A, i32 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x float> [[TMP1]], float 4.200000e+01, i32 1
-; CHECK-NEXT:    ret <2 x float> [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x float> <float undef, float 4.200000e+01>, float %A, i32 0
+; CHECK-NEXT:    ret <2 x float> [[TMP1]]
 ;
   %tmp23 = bitcast float %A to i32
   %tmp24 = zext i32 %tmp23 to i64
   %tmp25 = shl i64 %tmp24, 32
   %mask20 = or i64 %tmp25, 1109917696
   %tmp35 = bitcast i64 %mask20 to <2 x float>
   ret <2 x float> %tmp35
 }

test/Transforms/InstCombine/insert-extract-shuffle.ll

;
   ret <8 x float> %i1
 }
 
 ; PR26015: https://llvm.org/bugs/show_bug.cgi?id=26015
 ; The widening shuffle must be inserted before any uses.
 
 define <8 x i16> @pr26015(<4 x i16> %t0) {
 ; CHECK-LABEL: @pr26015(
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i16> %t0, <4 x i16> undef, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
-; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <4 x i16> %t0, i32 2
-; CHECK-NEXT:    [[T2:%.*]] = insertelement <8 x i16> <i16 0, i16 0, i16 0, i16 undef, i16 0, i16 0, i16 undef, i16 undef>, i16 [[TMP2]], i32 3
-; CHECK-NEXT:    [[T3:%.*]] = insertelement <8 x i16> [[T2]], i16 0, i32 6
-; CHECK-NEXT:    [[T5:%.*]] = shufflevector <8 x i16> [[T3]], <8 x i16> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 11>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i16> %t0, <4 x i16> undef, <8 x i32> <i32 undef, i32 undef, i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[T5:%.*]] = shufflevector <8 x i16> <i16 0, i16 0, i16 0, i16 undef, i16 0, i16 0, i16 0, i16 undef>, <8 x i16> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 10, i32 4, i32 5, i32 6, i32 11>
 ; CHECK-NEXT:    ret <8 x i16> [[T5]]
 ;
   %t1 = extractelement <4 x i16> %t0, i32 2
   %t2 = insertelement <8 x i16> zeroinitializer, i16 %t1, i32 3
   %t3 = insertelement <8 x i16> %t2, i16 0, i32 6
   %t4 = extractelement <4 x i16> %t0, i32 3
   %t5 = insertelement <8 x i16> %t3, i16 %t4, i32 7
   ret <8 x i16> %t5
 }
 
 ; PR25999: https://llvm.org/bugs/show_bug.cgi?id=25999
 ; TODO: The widening shuffle could be inserted at the start of the function to allow the first extract to use it.
 
 define <8 x i16> @pr25999(<4 x i16> %t0, i1 %b) {
 ; CHECK-LABEL: @pr25999(
 ; CHECK-NEXT:    [[T1:%.*]] = extractelement <4 x i16> %t0, i32 2
 ; CHECK-NEXT:    br i1 %b, label %if, label %end
 ; CHECK:       if:
 ; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i16> %t0, <4 x i16> undef, <8 x i32> <i32 undef, i32 undef, i32 undef, i32 3, i32 undef, i32 undef, i32 undef, i32 undef>
-; CHECK-NEXT:    [[T2:%.*]] = insertelement <8 x i16> <i16 0, i16 0, i16 0, i16 undef, i16 0, i16 0, i16 undef, i16 undef>, i16 [[T1]], i32 3
-; CHECK-NEXT:    [[T3:%.*]] = insertelement <8 x i16> [[T2]], i16 0, i32 6
+; CHECK-NEXT:    [[T3:%.*]] = insertelement <8 x i16> <i16 0, i16 0, i16 0, i16 undef, i16 0, i16 0, i16 0, i16 undef>, i16 [[T1]], i32 3
 ; CHECK-NEXT:    [[T5:%.*]] = shufflevector <8 x i16> [[T3]], <8 x i16> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 11>
 ; CHECK-NEXT:    ret <8 x i16> [[T5]]
 ; CHECK:       end:
 ; CHECK-NEXT:    [[A1:%.*]] = add i16 [[T1]], 4
 ; CHECK-NEXT:    [[T6:%.*]] = insertelement <8 x i16> <i16 undef, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>, i16 [[A1]], i32 0
 ; CHECK-NEXT:    ret <8 x i16> [[T6]]
 ;
 

test/Transforms/InstCombine/vec_demanded_elts.ll

 ;
   %shuffle.i = shufflevector <2 x float> %y, <2 x float> %y, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %shuffle9.i = shufflevector <4 x float> %x, <4 x float> %shuffle.i, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x float> %shuffle9.i
 }
 
 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:    [[TMP1:%.*]] = insertelement <2 x double> <double undef, double 0.000000e+00>, double %f, i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = fptrunc <2 x double> [[TMP1]] to <2 x float>
+; CHECK-NEXT:    ret <2 x float> [[TMP2]]
 ;
   %tmp9 = insertelement <4 x double> undef, double %f, i32 0
   %tmp10 = insertelement <4 x double> %tmp9, double 0.000000e+00, i32 1
   %tmp11 = insertelement <4 x double> %tmp10, double 0.000000e+00, i32 2
   %tmp12 = insertelement <4 x double> %tmp11, double 0.000000e+00, i32 3
   %tmp5 = fptrunc <4 x double> %tmp12 to <4 x float>
   %ret = shufflevector <4 x float> %tmp5, <4 x float> undef, <2 x i32> <i32 0, i32 1>
   ret <2 x float> %ret
 }
 
 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:    [[TMP1:%.*]] = insertelement <2 x float> <float undef, float 0.000000e+00>, float %f, i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = fpext <2 x float> [[TMP1]] to <2 x double>
+; CHECK-NEXT:    ret <2 x double> [[TMP2]]
 ;
   %tmp9 = insertelement <4 x float> undef, float %f, i32 0
   %tmp10 = insertelement <4 x float> %tmp9, float 0.000000e+00, i32 1
   %tmp11 = insertelement <4 x float> %tmp10, float 0.000000e+00, i32 2
   %tmp12 = insertelement <4 x float> %tmp11, float 0.000000e+00, i32 3
   %tmp5 = fpext <4 x float> %tmp12 to <4 x double>
   %ret = shufflevector <4 x double> %tmp5, <4 x double> undef, <2 x i32> <i32 0, i32 1>
   ret <2 x double> %ret
 }
 
 define <4 x double> @test_shuffle(<4 x double> %f) {
 ; CHECK-LABEL: @test_shuffle(
 ; CHECK-NEXT:    [[RET1:%.*]] = insertelement <4 x double> %f, double 1.000000e+00, i32 3
 ; CHECK-NEXT:    ret <4 x double> [[RET1]]
 ;
   %ret = shufflevector <4 x double> %f, <4 x double> <double undef, double 1.0, double undef, double undef>, <4 x i32> <i32 0, i32 1, i32 2, i32 5>
   ret <4 x double> %ret
 }
 
 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:%.*]] = insertelement <4 x float> <float undef, float undef, float undef, float 3.000000e+00>, float %f, i32 0
 ; CHECK-NEXT:    [[RET:%.*]] = shufflevector <4 x float> [[A3]], <4 x float> <float undef, float 4.000000e+00, float 5.000000e+00, float undef>, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT:    ret <4 x float> [[RET]]
 ;
   %a0 = insertelement <4 x float> undef, float %f, i32 0
   %a1 = insertelement <4 x float> %a0, float 1.000000e+00, i32 1
   %a2 = insertelement <4 x float> %a1, float 2.000000e+00, i32 2
   %a3 = insertelement <4 x float> %a2, float 3.000000e+00, i32 3
   %b0 = insertelement <4 x float> undef, float %g, i32 0