[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

781

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

794

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

Instruction *InsertFirst = InsertElementInst::Create(
       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);
 }
 
+/// If we have an insertelement instruction feeding into another insertelement
+/// and the 2nd is inserting a constant into the vector, canonicalize that
+/// constant insertion before the insertion of a variable:
+///
+/// insertelement (insertelement X, Y, IdxC1), ScalarC, IdxC2 -->
+/// insertelement (insertelement X, ScalarC, IdxC2), Y, IdxC1
+///
+/// This has the potential of eliminating the 2nd insertelement instruction
+/// via constant folding of the scalar constant into a vector constant.
+static Instruction *hoistInsEltConst(InsertElementInst &InsElt2,
+                                     InstCombiner::BuilderTy &Builder) {
+  auto *InsElt1 = dyn_cast<InsertElementInst>(InsElt2.getOperand(0));
+  if (!InsElt1 || !InsElt1->hasOneUse())
+    return nullptr;
+
+  Value *X, *Y;
+  Constant *ScalarC;
+  ConstantInt *IdxC1, *IdxC2;
+  if (match(InsElt1->getOperand(0), m_Value(X)) &&
+      match(InsElt1->getOperand(1), m_Value(Y)) && !isa<Constant>(Y) &&
+      match(InsElt1->getOperand(2), m_ConstantInt(IdxC1)) &&
+      match(InsElt2.getOperand(1), m_Constant(ScalarC)) &&
+      match(InsElt2.getOperand(2), m_ConstantInt(IdxC2)) && IdxC1 != IdxC2) {
+    Value *NewInsElt1 = Builder.CreateInsertElement(X, ScalarC, IdxC2);
+    return InsertElementInst::Create(NewInsElt1, Y, IdxC1);
+  }
+
+  return nullptr;
+}
+
 /// 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;
static Instruction *foldConstantInsEltIntoShuffle(InsertElementInst &InsElt) {
   }
   return nullptr;
 }
 
 Instruction *InstCombiner::visitInsertElementInst(InsertElementInst &IE) {
   Value *VecOp    = IE.getOperand(0);
   Value *ScalarOp = IE.getOperand(1);
   Value *IdxOp    = IE.getOperand(2);
 
   // Inserting an undef or into an undefined place, remove this.
   if (isa<UndefValue>(ScalarOp) || isa<UndefValue>(IdxOp))
     replaceInstUsesWith(IE, VecOp);
 
   // If the inserted element was extracted from some other vector, and if the
   // indexes are constant, try to turn this into a shufflevector operation.
   if (ExtractElementInst *EI = dyn_cast<ExtractElementInst>(ScalarOp)) {
     if (isa<ConstantInt>(EI->getOperand(1)) && isa<ConstantInt>(IdxOp)) {
       unsigned NumInsertVectorElts = IE.getType()->getNumElements();
       unsigned NumExtractVectorElts =
           EI->getOperand(0)->getType()->getVectorNumElements();
       unsigned ExtractedIdx =
         cast<ConstantInt>(EI->getOperand(1))->getZExtValue();
       unsigned InsertedIdx = cast<ConstantInt>(IdxOp)->getZExtValue();
 
       if (ExtractedIdx >= NumExtractVectorElts) // Out of range extract.
         return replaceInstUsesWith(IE, VecOp);
 
       if (InsertedIdx >= NumInsertVectorElts)  // Out of range insert.
         return replaceInstUsesWith(IE, UndefValue::get(IE.getType()));
 
if (Value *V = SimplifyDemandedVectorElts(&IE, AllOnesEltMask, UndefElts)) {
     if (V != &IE)
       return replaceInstUsesWith(IE, V);
     return &IE;
   }
 
   if (Instruction *Shuf = foldConstantInsEltIntoShuffle(IE))
     return Shuf;
 
+  if (Instruction *NewInsElt = hoistInsEltConst(IE, *Builder))
+    return NewInsElt;
+
   // Turn a sequence of inserts that broadcasts a scalar into a single
   // insert + shufflevector.
   if (Instruction *Broadcast = foldInsSequenceIntoBroadcast(IE))
     return Broadcast;
 
   return nullptr;
 }
 

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

test/Transforms/InstCombine/vector_insertelt_shuffle.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; insertelements should fold to shuffle
 define <4 x float> @foo(<4 x float> %x) {
 ; CHECK-LABEL: @foo(
 ; CHECK-NEXT:    [[INS2:%.*]] = shufflevector <4 x float> %x, <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> [[INS2]]
 ;
   %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
 }
 
+; Insert of a constant is canonicalized ahead of insert of a variable.
+
 define <4 x float> @bar(<4 x float> %x, float %a) {
 ; CHECK-LABEL: @bar(
-; CHECK-NEXT:    [[INS1:%.*]] = insertelement <4 x float> %x, float %a, i32 1
-; CHECK-NEXT:    [[INS2:%.*]] = insertelement <4 x float> [[INS1]], float 2.000000e+00, i32 2
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x float> %x, float 2.000000e+00, i32 2
+; CHECK-NEXT:    [[INS2:%.*]] = insertelement <4 x float> [[TMP1]], float %a, i32 1
 ; CHECK-NEXT:    ret <4 x float> [[INS2]]
 ;
   %ins1 = insertelement<4 x float> %x, float %a, i32 1
   %ins2 = insertelement<4 x float> %ins1, float 2.0, i32 2
   ret <4 x float> %ins2
 }
 
 define <4 x float> @baz(<4 x float> %x, i32 %a) {