[InstCombine] move extend after insertelement if both operands are extended

[spatel]

Summary

I was wondering how instcombine does on the examples in D109236, and we're missing what looks like a basic transform:

inselt (ext X), (ext Y), Index --> ext (inselt X, Y, Index)

https://alive2.llvm.org/ce/z/z2aBu9

If this looks ok, then instcombine can handle the relatively simple tests in the other patch.
So we might want to make those a bit longer to show the power of the more general approach in -aggressive-instcombine.

[lebedev.ri]

llvm/test/Transforms/InstCombine/insert-ext.ll

62–64

I guess it will depend on our profitability/correctness reasoning for this transform,
but under standard instcombine profitability model we should be fine producing

%s = sext i8 %y to i16
%i = insertelement <2 x i16> %x, i16 %s, i32 %index
%r = sext <2 x i16> %i to <2 x i32>

[spatel]

llvm/test/Transforms/InstCombine/insert-ext.ll

62–64

Yes - I thought about including this one, but it seemed improbable and makes the use check a bit more restrictive. Ie, if we're creating 3 new insts, then we need one-use on both of the operands.

And I wasn't sure, but it does seem to hold for fpext too:
https://alive2.llvm.org/ce/z/QNQKWt

I can make that a TODO if that's ok.

[RKSimon]

Should we handle trunc as well? https://alive2.llvm.org/ce/z/93QN9g

[spatel]

In D109537#2992713, @RKSimon wrote:

Should we handle trunc as well? https://alive2.llvm.org/ce/z/93QN9g

I want to check that codegen isn't impacted by using insertelement on a wider (possible illegal type) vector op, but I think the answer is 'yes'.

In fact, we should probably do this for all cast opcodes to be consistent. I have to go through the whole opcode list (and add piles of tests) to be sure, but it works for these at least:
https://alive2.llvm.org/ce/z/HxH4AC (fptosi / fptoui / sitofp / uitofp)
https://alive2.llvm.org/ce/z/Rh_dqv (bitcast)

Another TODO for this patch, or do we prefer to handle it all in one patch?

[spatel]

In D109537#2994331, @spatel wrote:

https://alive2.llvm.org/ce/z/Rh_dqv (bitcast)

Just noticed that we already have this one at line 1436 of the existing source file:

// If the vector and scalar are both bitcast from the same element type, do
// the insert in that source type followed by bitcast.

[spatel]

In D109537#2994331, @spatel wrote:

In D109537#2992713, @RKSimon wrote:

Should we handle trunc as well? https://alive2.llvm.org/ce/z/93QN9g

I want to check that codegen isn't impacted by using insertelement on a wider (possible illegal type) vector op, but I think the answer is 'yes'.

We would need to deal with something like this in the backend to be safe:
https://godbolt.org/z/E1nKvd6e3

[RKSimon]

In D109537#2994425, @spatel wrote:

In D109537#2994331, @spatel wrote:

In D109537#2992713, @RKSimon wrote:

Should we handle trunc as well? https://alive2.llvm.org/ce/z/93QN9g

I want to check that codegen isn't impacted by using insertelement on a wider (possible illegal type) vector op, but I think the answer is 'yes'.

We would need to deal with something like this in the backend to be safe:
https://godbolt.org/z/E1nKvd6e3

Yeah, there's going to be a number of these - in particular for integer insertions

[anton-afanasyev]

If this looks ok, then instcombine can handle the relatively simple tests in the other patch.
So we might want to make those a bit longer to show the power of the more general approach in -aggressive-instcombine.

No one of tests touched by D109236 isn't processed by this patch, since they are not matched by inselt (ext X), (ext Y), Index pattern.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1430

Move hasOneUse() check above, before matching? Or should be here in light of TODO?

[spatel]

Patch updated:

1. Moved use check up.
2. Added TODO comments - up to 3 now :)

[spatel]

In D109537#2997664, @anton-afanasyev wrote:

If this looks ok, then instcombine can handle the relatively simple tests in the other patch.
So we might want to make those a bit longer to show the power of the more general approach in -aggressive-instcombine.

No one of tests touched by D109236 isn't processed by this patch, since they are not matched by inselt (ext X), (ext Y), Index pattern.

What I meant is that the tests in that patch are reduced by instcombine after this patch, so there's nothing left for AIC to do if regular instcombine runs first:

$ opt -instcombine 109236.ll -S

define <2 x i16> @extract_insert(<2 x i8> %a, <2 x i8> %b) {
  %1 = shufflevector <2 x i8> %b, <2 x i8> %a, <2 x i32> <i32 0, i32 2>
  %trunc = zext <2 x i8> %1 to <2 x i16>
  ret <2 x i16> %trunc
}

define <2 x i16> @insert_poison(i8 %a) {
  %1 = zext i8 %a to i16
  %trunc = insertelement <2 x i16> undef, i16 %1, i32 0
  ret <2 x i16> %trunc
}

[spatel]

Patch updated:
I missed moving the code comment above the hasOneUse() check in the previous draft.

[anton-afanasyev]

In D109537#3001567, @spatel wrote:

What I meant is that the tests in that patch are reduced by instcombine after this patch, so there's nothing left for AIC to do if regular instcombine runs first:

$ opt -instcombine 109236.ll -S

define <2 x i16> @extract_insert(<2 x i8> %a, <2 x i8> %b) {
  %1 = shufflevector <2 x i8> %b, <2 x i8> %a, <2 x i32> <i32 0, i32 2>
  %trunc = zext <2 x i8> %1 to <2 x i16>
  ret <2 x i16> %trunc
}

define <2 x i16> @insert_poison(i8 %a) {
  %1 = zext i8 %a to i16
  %trunc = insertelement <2 x i16> undef, i16 %1, i32 0
  ret <2 x i16> %trunc
}

Ok, I see. Rebased D109236 agains new test added.

[anton-afanasyev]

LGTM

llvm/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.
     return new ShuffleVectorInst(IEI->getOperand(0),
                                  ConstantVector::get(Values), Mask);
   }
   return nullptr;
 }
 
+/// If both the base vector and the inserted element are extended from the same
+/// type, do the insert element in the narrow source type followed by extend.
+/// TODO: This can be extended to include other cast opcodes, but particularly
+///       if we create a wider insertelement, make sure codegen is not harmed.
+static Instruction *narrowInsElt(InsertElementInst &InsElt,
+                                 InstCombiner::BuilderTy &Builder) {
+  // We are creating a vector extend. If the original vector extend has another
+  // use, that would mean we end up with 2 vector extends, so avoid that.
+  // TODO: We could ease the use-clause to "if at least one op has one use"
+  //       (assuming that the source types match - see next TODO comment).
+  Value *Vec = InsElt.getOperand(0);
+  if (!Vec->hasOneUse())
+    return nullptr;
+
+  Value *Scalar = InsElt.getOperand(1);
+  Value *X, *Y;
+  CastInst::CastOps CastOpcode;
+  if (match(Vec, m_FPExt(m_Value(X))) && match(Scalar, m_FPExt(m_Value(Y))))
+    CastOpcode = Instruction::FPExt;
+  else if (match(Vec, m_SExt(m_Value(X))) && match(Scalar, m_SExt(m_Value(Y))))
+    CastOpcode = Instruction::SExt;
+  else if (match(Vec, m_ZExt(m_Value(X))) && match(Scalar, m_ZExt(m_Value(Y))))
+    CastOpcode = Instruction::ZExt;
+  else
+    return nullptr;
+
+  // TODO: We can allow mismatched types by creating an intermediate cast.
+  if (X->getType()->getScalarType() != Y->getType())
+    return nullptr;
+
+  // inselt (ext X), (ext Y), Index --> ext (inselt X, Y, Index)
+  Value *NewInsElt = Builder.CreateInsertElement(X, Y, InsElt.getOperand(2));
+  return CastInst::Create(CastOpcode, NewInsElt, InsElt.getType());
+}
+
 Instruction *InstCombinerImpl::visitInsertElementInst(InsertElementInst &IE) {
   Value *VecOp    = IE.getOperand(0);
   Value *ScalarOp = IE.getOperand(1);
   Value *IdxOp    = IE.getOperand(2);
 
   if (auto *V = SimplifyInsertElementInst(
           VecOp, ScalarOp, IdxOp, SQ.getWithInstruction(&IE)))
     return replaceInstUsesWith(IE, V);
if (Instruction *Broadcast = foldInsSequenceIntoSplat(IE))
     return Broadcast;
 
   if (Instruction *Splat = foldInsEltIntoSplat(IE))
     return Splat;
 
   if (Instruction *IdentityShuf = foldInsEltIntoIdentityShuffle(IE))
     return IdentityShuf;
 
+  if (Instruction *Ext = narrowInsElt(IE, Builder))
+    return Ext;
+
   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/test/Transforms/InstCombine/insert-ext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 declare void @use(i32)
 declare void @usevec(<2 x i32>)
 
 define <2 x double> @fpext_fpext(<2 x half> %x, half %y, i32 %index) {
 ; CHECK-LABEL: @fpext_fpext(
-; CHECK-NEXT:    [[V:%.*]] = fpext <2 x half> [[X:%.*]] to <2 x double>
-; CHECK-NEXT:    [[S:%.*]] = fpext half [[Y:%.*]] to double
-; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x double> [[V]], double [[S]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x half> [[X:%.*]], half [[Y:%.*]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = fpext <2 x half> [[TMP1]] to <2 x double>
 ; CHECK-NEXT:    ret <2 x double> [[I]]
 ;
   %v = fpext <2 x half> %x to <2 x double>
   %s = fpext half %y to double
   %i = insertelement <2 x double> %v, double %s, i32 %index
   ret <2 x double> %i
 }
 
 define <2 x i32> @sext_sext(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @sext_sext(
-; CHECK-NEXT:    [[V:%.*]] = sext <2 x i8> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i32> [[V]], i32 [[S]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i8> [[X:%.*]], i8 [[Y:%.*]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = sext <2 x i8> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[I]]
 ;
   %v = sext <2 x i8> %x to <2 x i32>
   %s = sext i8 %y to i32
   %i = insertelement <2 x i32> %v, i32 %s, i32 %index
   ret <2 x i32> %i
 }
 
 define <2 x i12> @zext_zext(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @zext_zext(
-; CHECK-NEXT:    [[V:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i12>
-; CHECK-NEXT:    [[S:%.*]] = zext i8 [[Y:%.*]] to i12
-; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i12> [[V]], i12 [[S]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i8> [[X:%.*]], i8 [[Y:%.*]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = zext <2 x i8> [[TMP1]] to <2 x i12>
 ; CHECK-NEXT:    ret <2 x i12> [[I]]
 ;
   %v = zext <2 x i8> %x to <2 x i12>
   %s = zext i8 %y to i12
   %i = insertelement <2 x i12> %v, i12 %s, i32 %index
   ret <2 x i12> %i
 }
 
+; negative test - need same source type
+
 define <2 x double> @fpext_fpext_types(<2 x half> %x, float %y, i32 %index) {
 ; CHECK-LABEL: @fpext_fpext_types(
 ; CHECK-NEXT:    [[V:%.*]] = fpext <2 x half> [[X:%.*]] to <2 x double>
 ; CHECK-NEXT:    [[S:%.*]] = fpext float [[Y:%.*]] to double
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x double> [[V]], double [[S]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret <2 x double> [[I]]
 ;
   %v = fpext <2 x half> %x to <2 x double>
   %s = fpext float %y to double
   %i = insertelement <2 x double> %v, double %s, i32 %index
   ret <2 x double> %i
 }
 
+; negative test - need same source type
+
 define <2 x i32> @sext_sext_types(<2 x i16> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @sext_sext_types(
 ; CHECK-NEXT:    [[V:%.*]] = sext <2 x i16> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = sext i8 [[Y:%.*]] to i32
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i32> [[V]], i32 [[S]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[I]]
 ;
   %v = sext <2 x i16> %x to <2 x i32>
   %s = sext i8 %y to i32
   %i = insertelement <2 x i32> %v, i32 %s, i32 %index
   ret <2 x i32> %i
 }
 
+; negative test - need same extend opcode
+
 define <2 x i12> @sext_zext(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @sext_zext(
 ; CHECK-NEXT:    [[V:%.*]] = sext <2 x i8> [[X:%.*]] to <2 x i12>
 ; CHECK-NEXT:    [[S:%.*]] = zext i8 [[Y:%.*]] to i12
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i12> [[V]], i12 [[S]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret <2 x i12> [[I]]
 ;
   %v = sext <2 x i8> %x to <2 x i12>
   %s = zext i8 %y to i12
   %i = insertelement <2 x i12> %v, i12 %s, i32 %index
   ret <2 x i12> %i
 }
 
+; negative test - don't trade scalar extend for vector extend
+
 define <2 x i32> @sext_sext_use1(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @sext_sext_use1(
 ; CHECK-NEXT:    [[V:%.*]] = sext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    call void @usevec(<2 x i32> [[V]])
 ; CHECK-NEXT:    [[S:%.*]] = sext i8 [[Y:%.*]] to i32
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i32> [[V]], i32 [[S]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[I]]
 ;
   %v = sext <2 x i8> %x to <2 x i32>
   call void @usevec(<2 x i32> %v)
   %s = sext i8 %y to i32
   %i = insertelement <2 x i32> %v, i32 %s, i32 %index
   ret <2 x i32> %i
 }
 
 define <2 x i32> @zext_zext_use2(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @zext_zext_use2(
-; CHECK-NEXT:    [[V:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = zext i8 [[Y:%.*]] to i32
 ; CHECK-NEXT:    call void @use(i32 [[S]])
-; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i32> [[V]], i32 [[S]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i8> [[X:%.*]], i8 [[Y]], i32 [[INDEX:%.*]]
+; CHECK-NEXT:    [[I:%.*]] = zext <2 x i8> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[I]]
 ;
   %v = zext <2 x i8> %x to <2 x i32>
   %s = zext i8 %y to i32
   call void @use(i32 %s)
   %i = insertelement <2 x i32> %v, i32 %s, i32 %index
   ret <2 x i32> %i
 }
 
+; negative test - don't create an extra extend
+
 define <2 x i32> @zext_zext_use3(<2 x i8> %x, i8 %y, i32 %index) {
 ; CHECK-LABEL: @zext_zext_use3(
 ; CHECK-NEXT:    [[V:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    call void @usevec(<2 x i32> [[V]])
 ; CHECK-NEXT:    [[S:%.*]] = zext i8 [[Y:%.*]] to i32
 ; CHECK-NEXT:    call void @use(i32 [[S]])
 ; CHECK-NEXT:    [[I:%.*]] = insertelement <2 x i32> [[V]], i32 [[S]], i32 [[INDEX:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[I]]
 ;
   %v = zext <2 x i8> %x to <2 x i32>
   call void @usevec(<2 x i32> %v)
   %s = zext i8 %y to i32
   call void @use(i32 %s)
   %i = insertelement <2 x i32> %v, i32 %s, i32 %index
   ret <2 x i32> %i
 }