[InstCombine] Extend pattern to replace shuffle's insertelement operand

[piotr]

Summary

In D71220 a pattern was added to replace shuffle's insertelement operand
if inserted scalar is not demanded. The pattern was added only for
the case where the shuffle's mask size is equal to element's vector size.
However, that condition is not required because the pattern does not
change the shuffle vector size.

This patch extends the pattern to also include cases where shuffle's mask
size is not equal to element's vector size.

[piotr]

Tests pre-committed in 7457fe3dd44a0bc4b0296149c48188accefda5fa.

[foad]

Looks reasonable to me.

[piotr]

Ping.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

2294

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

 /// Try to replace a shuffle with an insertelement or try to replace a shuffle
 /// operand with the operand of an insertelement.
 static Instruction *foldShuffleWithInsert(ShuffleVectorInst &Shuf,
                                           InstCombinerImpl &IC) {
   Value *V0 = Shuf.getOperand(0), *V1 = Shuf.getOperand(1);
   SmallVector<int, 16> Mask;
   Shuf.getShuffleMask(Mask);
 
-  // The shuffle must not change vector sizes.
-  // TODO: This restriction could be removed if the insert has only one use
-  //       (because the transform would require a new length-changing shuffle).
   int NumElts = Mask.size();
-  if (NumElts != (int)(cast<FixedVectorType>(V0->getType())->getNumElements()))
-    return nullptr;
+  int InpNumElts = cast<FixedVectorType>(V0->getType())->getNumElements();
 
   // This is a specialization of a fold in SimplifyDemandedVectorElts. We may
   // not be able to handle it there if the insertelement has >1 use.
   // If the shuffle has an insertelement operand but does not choose the
   // inserted scalar element from that value, then we can replace that shuffle
   // operand with the source vector of the insertelement.
   Value *X;
   uint64_t IdxC;
   if (match(V0, m_InsertElt(m_Value(X), m_Value(), m_ConstantInt(IdxC)))) {
     // shuf (inselt X, ?, IdxC), ?, Mask --> shuf X, ?, Mask
     if (!is_contained(Mask, (int)IdxC))
       return IC.replaceOperand(Shuf, 0, X);
   }
   if (match(V1, m_InsertElt(m_Value(X), m_Value(), m_ConstantInt(IdxC)))) {
     // Offset the index constant by the vector width because we are checking for
     // accesses to the 2nd vector input of the shuffle.
-    IdxC += NumElts;
+    IdxC += InpNumElts;
     // shuf ?, (inselt X, ?, IdxC), Mask --> shuf ?, X, Mask
     if (!is_contained(Mask, (int)IdxC))
       return IC.replaceOperand(Shuf, 1, X);
   }
+  // For the rest of the transform, the shuffle must not change vector sizes.
+  // TODO: This restriction could be removed if the insert has only one use
+  //       (because the transform would require a new length-changing shuffle).
+  if (NumElts != InpNumElts)
+    return nullptr;
 
   // shuffle (insert ?, Scalar, IndexC), V1, Mask --> insert V1, Scalar, IndexC'
   auto isShufflingScalarIntoOp1 = [&](Value *&Scalar, ConstantInt *&IndexC) {
     // We need an insertelement with a constant index.
     if (!match(V0, m_InsertElt(m_Value(), m_Value(Scalar),
                                m_ConstantInt(IndexC))))
       return false;
 

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

;
   %s = shufflevector <4 x float> %y, <4 x float> %ins, <4 x i32> <i32 3, i32 2, i32 1, i32 4>
   ret <4 x float> %s
 }
 
 define <5 x float> @insert_undemanded_element_unequal_length_op0(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_undemanded_element_unequal_length_op0(
 ; CHECK-NEXT:    [[INS:%.*]] = insertelement <4 x float> [[X:%.*]], float 4.200000e+01, i32 3
 ; CHECK-NEXT:    call void @use(<4 x float> [[INS]])
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[INS]], <4 x float> [[Y:%.*]], <5 x i32> <i32 undef, i32 0, i32 7, i32 1, i32 4>
+; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[X]], <4 x float> [[Y:%.*]], <5 x i32> <i32 undef, i32 0, i32 7, i32 1, i32 4>
 ; CHECK-NEXT:    ret <5 x float> [[S]]
 ;
   %ins = insertelement <4 x float> %x, float 42.0, i32 3
   call void @use(<4 x float> %ins)
   %s = shufflevector <4 x float> %ins, <4 x float> %y, <5 x i32> <i32 undef, i32 0, i32 7, i32 1, i32 4>
   ret <5 x float> %s
 }
 
 define <5 x float> @insert_undemanded_element_unequal_length_op1(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_undemanded_element_unequal_length_op1(
 ; CHECK-NEXT:    [[INS:%.*]] = insertelement <4 x float> [[X:%.*]], float 4.200000e+01, i32 3
 ; CHECK-NEXT:    call void @use(<4 x float> [[INS]])
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[Y:%.*]], <4 x float> [[INS]], <5 x i32> <i32 undef, i32 3, i32 2, i32 1, i32 4>
+; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[Y:%.*]], <4 x float> [[X]], <5 x i32> <i32 undef, i32 3, i32 2, i32 1, i32 4>
 ; CHECK-NEXT:    ret <5 x float> [[S]]
 ;
   %ins = insertelement <4 x float> %x, float 42.0, i32 3
   call void @use(<4 x float> %ins)
   %s = shufflevector <4 x float> %y, <4 x float> %ins, <5 x i32> <i32 undef, i32 3, i32 2, i32 1, i32 4>
   ret <5 x float> %s
 }