[InstCombine] replace shuffle's insertelement operand if inserted scalar is not demanded

[spatel]

Summary

This pattern is noted as a regression from:
D70246
...where we removed an over-aggressive shuffle simplification.

SimplifyDemandedVectorElts fails to catch this case when the insert has multiple uses, so I'm proposing to pattern match the minimal sequence directly. This fold does not conflict with any of our current shuffle undef/poison semantics AFAICT.

[lebedev.ri]

This is legal.

----------------------------------------
define <4 x float> @insert_undemanded_element_op0(<4 x float> %x, <4 x float> %y) {
%0:
  %ins = insertelement <4 x float> %x, float 42.000000, i32 3
  %s = shufflevector <4 x float> %ins, <4 x float> %y, <4 x i32> { 0, 7, 1, 4 }
  ret <4 x float> %s
}
=>
define <4 x float> @insert_undemanded_element_op0(<4 x float> %x, <4 x float> %y) {
%0:
  %ins = insertelement <4 x float> %x, float 42.000000, i32 3
  %s = shufflevector <4 x float> %x, <4 x float> %y, <4 x i32> { 0, 7, 1, 4 }
  ret <4 x float> %s
}
Transformation seems to be correct!

Summary:
  1 correct transformations
  0 incorrect transformations
  0 errors

[foad]

I can confirm that this fixes the regression I noted. Thanks!

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

static Instruction *foldIdentityExtractShuffle(ShuffleVectorInst &Shuf) {
   for (unsigned i = 0; i != NumElts; ++i) {
     Constant *ExtractMaskElt = Shuf.getMask()->getAggregateElement(i);
     Constant *MaskElt = Mask->getAggregateElement(i);
     NewMask[i] = isa<UndefValue>(ExtractMaskElt) ? ExtractMaskElt : MaskElt;
   }
   return new ShuffleVectorInst(X, Y, ConstantVector::get(NewMask));
 }
 
-/// Try to replace a shuffle with an insertelement.
+/// 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) {
   Value *V0 = Shuf.getOperand(0), *V1 = Shuf.getOperand(1);
   SmallVector<int, 16> Mask = Shuf.getShuffleMask();
 
   // 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)(V0->getType()->getVectorNumElements()))
     return nullptr;
 
+  // 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_InsertElement(m_Value(X), m_Value(), m_ConstantInt(IdxC)))) {
+    // shuf (inselt X, ?, IdxC), ?, Mask --> shuf X, ?, Mask
+    if (none_of(Mask, [IdxC](int MaskElt) { return MaskElt == (int)IdxC; })) {
+      Shuf.setOperand(0, X);
+      return &Shuf;
+    }
+  }
+  if (match(V1, m_InsertElement(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;
+    // shuf ?, (inselt X, ?, IdxC), Mask --> shuf ?, X, Mask
+    if (none_of(Mask, [IdxC](int MaskElt) { return MaskElt == (int)IdxC; })) {
+      Shuf.setOperand(1, X);
+      return &Shuf;
+    }
+  }
+
   // 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_InsertElement(m_Value(), m_Value(Scalar),
                                    m_ConstantInt(IndexC))))
       return false;
 
     // Test the shuffle mask to see if it splices the inserted scalar into the

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

;
   %v2 = insertelement <4 x i8> %v1, i8 %y, i32 3
   ret <4 x i8> %v2
 }
 
 define <4 x float> @insert_undemanded_element_op0(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_undemanded_element_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:%.*]], <4 x i32> <i32 0, i32 7, i32 1, i32 4>
+; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[X]], <4 x float> [[Y:%.*]], <4 x i32> <i32 0, i32 7, i32 1, i32 4>
 ; CHECK-NEXT:    ret <4 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, <4 x i32> <i32 0, i32 7, i32 1, i32 4>
   ret <4 x float> %s
 }
 
 define <4 x float> @insert_undemanded_element_op1(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_undemanded_element_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]], <4 x i32> <i32 3, i32 2, i32 1, i32 4>
+; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[Y:%.*]], <4 x float> [[X]], <4 x i32> <i32 3, i32 2, i32 1, i32 4>
 ; CHECK-NEXT:    ret <4 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, <4 x i32> <i32 3, i32 2, i32 1, i32 4>
   ret <4 x float> %s
 }
 
+; Negative test - shuffle chooses the inserted constant.
+
 define <4 x float> @insert_demanded_element_op0(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_demanded_element_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:%.*]], <4 x i32> <i32 3, i32 2, i32 1, i32 4>
 ; CHECK-NEXT:    ret <4 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, <4 x i32> <i32 3, i32 2, i32 1, i32 4>
   ret <4 x float> %s
 }
 
+; Negative test - shuffle chooses the inserted constant.
+
 define <4 x float> @insert_demanded_element_op1(<4 x float> %x, <4 x float> %y) {
 ; CHECK-LABEL: @insert_demanded_element_op1(
 ; CHECK-NEXT:    [[INS:%.*]] = insertelement <4 x float> [[X:%.*]], float 4.300000e+01, i32 3
 ; CHECK-NEXT:    call void @use(<4 x float> [[INS]])
 ; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x float> [[Y:%.*]], <4 x float> [[INS]], <4 x i32> <i32 0, i32 7, i32 1, i32 4>
 ; CHECK-NEXT:    ret <4 x float> [[S]]
 ;
   %ins = insertelement <4 x float> %x, float 43.0, i32 3
   call void @use(<4 x float> %ins)
   %s = shufflevector <4 x float> %y, <4 x float> %ins, <4 x i32> <i32 0, i32 7, i32 1, i32 4>
   ret <4 x float> %s
 }