[InstCombine] try to convert vector insert+extract to trunc

[spatel]

Summary

This transform is requested for the backend in:
https://bugs.llvm.org/show_bug.cgi?id=39016
...but I figured it was worth doing in IR too, and it's probably easier to implement here, so that's this patch.

In the simplest case, we are just truncating a scalar value. If the extract index doesn't correspond to the LSBs of the scalar, then we have to shift-right before the truncate. Endian-ness makes this tricky, but hopefully the ASCII-art helps visualize the transform.

[RKSimon]

Worth testing with more exotic types? <2 x i24> to <6 x i8> ?

[spatel]

In D52439#1246442, @RKSimon wrote:

Worth testing with more exotic types? <2 x i24> to <6 x i8> ?

Sure - but we can get even crazier than that. :)
Will update shortly...

[spatel]

Patch updated - no code changes, but:

1. Added crazy type example (3 x i15).
2. I committed the negative/partial negative tests along with the new test in rL343128.

[RKSimon]

LGTM

[spatel]

Reverted at rL343458 because we can create bogus trunc insts (not sure how yet).

[spatel]

In D52439#1250933, @spatel wrote:

Reverted at rL343458 because we can create bogus trunc insts (not sure how yet).

I forgot to add a check for casting to an FP type in the original patch, so we died trying to perform an integer op on a float.
For now, I just added that check, but I think we should be able to handle patterns with FP types too, so I added a TODO comment.

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

Instruction *InstCombiner::scalarizePHI(ExtractElementInst &EI, PHINode *PN) {
 
   for (auto E : Extracts)
     replaceInstUsesWith(*E, scalarPHI);
 
   return &EI;
 }
 
 static Instruction *foldBitcastExtElt(ExtractElementInst &Ext,
-                                      InstCombiner::BuilderTy &Builder) {
+                                      InstCombiner::BuilderTy &Builder,
+                                      bool IsBigEndian) {
   Value *X;
   uint64_t ExtIndexC;
   if (!match(Ext.getVectorOperand(), m_BitCast(m_Value(X))) ||
       !X->getType()->isVectorTy() ||
       !match(Ext.getIndexOperand(), m_ConstantInt(ExtIndexC)))
     return nullptr;
 
   // If this extractelement is using a bitcast from a vector of the same number
   // of elements, see if we can find the source element from the source vector:
   // extelt (bitcast VecX), IndexC --> bitcast X[IndexC]
   Type *SrcTy = X->getType();
   Type *DestTy = Ext.getType();
   unsigned NumSrcElts = SrcTy->getVectorNumElements();
   unsigned NumElts = Ext.getVectorOperandType()->getNumElements();
   if (NumSrcElts == NumElts)
     if (Value *Elt = findScalarElement(X, ExtIndexC))
       return new BitCastInst(Elt, DestTy);
 
+  // If the source elements are wider than the destination, try to shift and
+  // truncate a subset of scalar bits of an insert op.
+  if (NumSrcElts < NumElts && SrcTy->getScalarType()->isIntegerTy()) {
+    Value *Scalar;
+    uint64_t InsIndexC;
+    if (!match(X, m_InsertElement(m_Value(), m_Value(Scalar),
+                                  m_ConstantInt(InsIndexC))))
+      return nullptr;
+
+    // The extract must be from the subset of vector elements that we inserted
+    // into. Example: if we inserted element 1 of a <2 x i64> and we are
+    // extracting an i16 (narrowing ratio = 4), then this extract must be from 1
+    // of elements 4-7 of the bitcasted vector.
+    unsigned NarrowingRatio = NumElts / NumSrcElts;
+    if (ExtIndexC / NarrowingRatio != InsIndexC)
+      return nullptr;
+
+    // We are extracting part of the original scalar. How that scalar is
+    // inserted into the vector depends on the endian-ness. Example:
+    //              Vector Byte Elt Index:    0  1  2  3  4  5  6  7
+    //                                       +--+--+--+--+--+--+--+--+
+    // inselt <2 x i32> V, <i32> S, 1:       |V0|V1|V2|V3|S0|S1|S2|S3|
+    // extelt <4 x i16> V', 3:               |                 |S2|S3|
+    //                                       +--+--+--+--+--+--+--+--+
+    // If this is little-endian, S2|S3 are the MSB of the 32-bit 'S' value.
+    // If this is big-endian, S2|S3 are the LSB of the 32-bit 'S' value.
+    // In this example, we must right-shift little-endian. Big-endian is just a
+    // truncate.
+    unsigned Chunk = ExtIndexC % NarrowingRatio;
+    if (IsBigEndian)
+      Chunk = NarrowingRatio - 1 - Chunk;
+    unsigned ShAmt = Chunk * DestTy->getPrimitiveSizeInBits();
+    if (ShAmt) {
+      // Bail out if we could end with more instructions than we started with.
+      if (!Ext.getVectorOperand()->hasOneUse())
+        return nullptr;
+      Scalar = Builder.CreateLShr(Scalar, ShAmt);
+    }
+    return new TruncInst(Scalar, DestTy);
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombiner::visitExtractElementInst(ExtractElementInst &EI) {
   if (Value *V = SimplifyExtractElementInst(EI.getVectorOperand(),
                                             EI.getIndexOperand(),
                                             SQ.getWithInstruction(&EI)))
     return replaceInstUsesWith(EI, V);
if (EI.getOperand(0)->hasOneUse() && NumElts != 1) {
       DemandedMask.setBit(IdxC->getZExtValue());
       if (Value *V = SimplifyDemandedVectorElts(EI.getOperand(0), DemandedMask,
                                                 UndefElts)) {
         EI.setOperand(0, V);
         return &EI;
       }
     }
 
-    if (Instruction *I = foldBitcastExtElt(EI, Builder))
+    if (Instruction *I = foldBitcastExtElt(EI, Builder, DL.isBigEndian()))
       return I;
 
     // If there's a vector PHI feeding a scalar use through this extractelement
     // instruction, try to scalarize the PHI.
     if (PHINode *PN = dyn_cast<PHINode>(EI.getOperand(0))) {
       Instruction *scalarPHI = scalarizePHI(EI, PN);
       if (scalarPHI)
         return scalarPHI;

test/Transforms/InstCombine/extractelement.ll

;
   %vec = insertelement <8 x i64> undef, i64 %in, i32 0
   %splat = shufflevector <8 x i64> %vec, <8 x i64> undef, <8 x i32> zeroinitializer
   %add = add <8 x i64> %splat, <i64 0, i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7>
   %r = extractelement <8 x i64> %add, i32 0
   ret i64 %r
 }
 
 define i32 @bitcasted_inselt_wide_source_zero_elt(i64 %x) {
-; ANY-LABEL: @bitcasted_inselt_wide_source_zero_elt(
-; ANY-NEXT:    [[I:%.*]] = insertelement <2 x i64> undef, i64 [[X:%.*]], i32 0
-; ANY-NEXT:    [[B:%.*]] = bitcast <2 x i64> [[I]] to <4 x i32>
-; ANY-NEXT:    [[R:%.*]] = extractelement <4 x i32> [[B]], i32 0
-; ANY-NEXT:    ret i32 [[R]]
+; LE-LABEL: @bitcasted_inselt_wide_source_zero_elt(
+; LE-NEXT:    [[R:%.*]] = trunc i64 [[X:%.*]] to i32
+; LE-NEXT:    ret i32 [[R]]
+;
+; BE-LABEL: @bitcasted_inselt_wide_source_zero_elt(
+; BE-NEXT:    [[TMP1:%.*]] = lshr i64 [[X:%.*]], 32
+; BE-NEXT:    [[R:%.*]] = trunc i64 [[TMP1]] to i32
+; BE-NEXT:    ret i32 [[R]]
 ;
   %i = insertelement <2 x i64> zeroinitializer, i64 %x, i32 0
   %b = bitcast <2 x i64> %i to <4 x i32>
   %r = extractelement <4 x i32> %b, i32 0
   ret i32 %r
 }
 
 define i16 @bitcasted_inselt_wide_source_modulo_elt(i64 %x) {
-; ANY-LABEL: @bitcasted_inselt_wide_source_modulo_elt(
-; ANY-NEXT:    [[I:%.*]] = insertelement <2 x i64> undef, i64 [[X:%.*]], i32 1
-; ANY-NEXT:    [[B:%.*]] = bitcast <2 x i64> [[I]] to <8 x i16>
-; ANY-NEXT:    [[R:%.*]] = extractelement <8 x i16> [[B]], i32 4
-; ANY-NEXT:    ret i16 [[R]]
+; LE-LABEL: @bitcasted_inselt_wide_source_modulo_elt(
+; LE-NEXT:    [[R:%.*]] = trunc i64 [[X:%.*]] to i16
+; LE-NEXT:    ret i16 [[R]]
+;
+; BE-LABEL: @bitcasted_inselt_wide_source_modulo_elt(
+; BE-NEXT:    [[TMP1:%.*]] = lshr i64 [[X:%.*]], 48
+; BE-NEXT:    [[R:%.*]] = trunc i64 [[TMP1]] to i16
+; BE-NEXT:    ret i16 [[R]]
 ;
   %i = insertelement <2 x i64> undef, i64 %x, i32 1
   %b = bitcast <2 x i64> %i to <8 x i16>
   %r = extractelement <8 x i16> %b, i32 4
   ret i16 %r
 }
 
 define i32 @bitcasted_inselt_wide_source_not_modulo_elt(i64 %x) {
-; ANY-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt(
-; ANY-NEXT:    [[I:%.*]] = insertelement <2 x i64> undef, i64 [[X:%.*]], i32 0
-; ANY-NEXT:    [[B:%.*]] = bitcast <2 x i64> [[I]] to <4 x i32>
-; ANY-NEXT:    [[R:%.*]] = extractelement <4 x i32> [[B]], i32 1
-; ANY-NEXT:    ret i32 [[R]]
+; LE-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt(
+; LE-NEXT:    [[TMP1:%.*]] = lshr i64 [[X:%.*]], 32
+; LE-NEXT:    [[R:%.*]] = trunc i64 [[TMP1]] to i32
+; LE-NEXT:    ret i32 [[R]]
+;
+; BE-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt(
+; BE-NEXT:    [[R:%.*]] = trunc i64 [[X:%.*]] to i32
+; BE-NEXT:    ret i32 [[R]]
 ;
   %i = insertelement <2 x i64> undef, i64 %x, i32 0
   %b = bitcast <2 x i64> %i to <4 x i32>
   %r = extractelement <4 x i32> %b, i32 1
   ret i32 %r
 }
 
 define i8 @bitcasted_inselt_wide_source_not_modulo_elt_not_half(i32 %x) {
-; ANY-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt_not_half(
-; ANY-NEXT:    [[I:%.*]] = insertelement <2 x i32> undef, i32 [[X:%.*]], i32 0
-; ANY-NEXT:    [[B:%.*]] = bitcast <2 x i32> [[I]] to <8 x i8>
+; LE-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt_not_half(
+; LE-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 16
+; LE-NEXT:    [[R:%.*]] = trunc i32 [[TMP1]] to i8
+; LE-NEXT:    ret i8 [[R]]
+;
+; BE-LABEL: @bitcasted_inselt_wide_source_not_modulo_elt_not_half(
+; BE-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 8
+; BE-NEXT:    [[R:%.*]] = trunc i32 [[TMP1]] to i8
+; BE-NEXT:    ret i8 [[R]]
+;
+  %i = insertelement <2 x i32> undef, i32 %x, i32 0
+  %b = bitcast <2 x i32> %i to <8 x i8>
+  %r = extractelement <8 x i8> %b, i32 2
+  ret i8 %r
+}
+
+; Negative test for the above fold, but we can remove the insert here.
+
+define i8 @bitcasted_inselt_wide_source_wrong_insert(<2 x i32> %v, i32 %x) {
+; ANY-LABEL: @bitcasted_inselt_wide_source_wrong_insert(
+; ANY-NEXT:    [[B:%.*]] = bitcast <2 x i32> [[V:%.*]] to <8 x i8>
 ; ANY-NEXT:    [[R:%.*]] = extractelement <8 x i8> [[B]], i32 2
 ; ANY-NEXT:    ret i8 [[R]]
 ;
-  %i = insertelement <2 x i32> undef, i32 %x, i32 0
+  %i = insertelement <2 x i32> %v, i32 %x, i32 1
   %b = bitcast <2 x i32> %i to <8 x i8>
   %r = extractelement <8 x i8> %b, i32 2
   ret i8 %r
 }
 
+; Partial negative test for the above fold, extra uses are not allowed if shift is needed.
+
+declare void @use(<8 x i8>)
+
+define i8 @bitcasted_inselt_wide_source_uses(i32 %x) {
+; LE-LABEL: @bitcasted_inselt_wide_source_uses(
+; LE-NEXT:    [[I:%.*]] = insertelement <2 x i32> undef, i32 [[X:%.*]], i32 0
+; LE-NEXT:    [[B:%.*]] = bitcast <2 x i32> [[I]] to <8 x i8>
+; LE-NEXT:    call void @use(<8 x i8> [[B]])
+; LE-NEXT:    [[R:%.*]] = extractelement <8 x i8> [[B]], i32 3
+; LE-NEXT:    ret i8 [[R]]
+;
+; BE-LABEL: @bitcasted_inselt_wide_source_uses(
+; BE-NEXT:    [[I:%.*]] = insertelement <2 x i32> undef, i32 [[X:%.*]], i32 0
+; BE-NEXT:    [[B:%.*]] = bitcast <2 x i32> [[I]] to <8 x i8>
+; BE-NEXT:    call void @use(<8 x i8> [[B]])
+; BE-NEXT:    [[R:%.*]] = trunc i32 [[X]] to i8
+; BE-NEXT:    ret i8 [[R]]
+;
+  %i = insertelement <2 x i32> undef, i32 %x, i32 0
+  %b = bitcast <2 x i32> %i to <8 x i8>
+  call void @use(<8 x i8> %b)
+  %r = extractelement <8 x i8> %b, i32 3
+  ret i8 %r
+}
+