[InstSimplify] fold splat of inserted constant to vector constant

[spatel]

Summary

shuf (inselt ?, C, IndexC), undef, <IndexC, IndexC...> --> <C, C...>

This is another missing shuffle fold pattern uncovered by the shuffle correctness fix from D70246.
The problem was visible in the post-commit thread example posted by @foad, but we managed to overcome the limitation for that particular case with D71220.
This is something like the inverse of the previous fix - there we didn't demand the inserted scalar, and here we are only demanding an inserted scalar.

[lebedev.ri]

LG

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyShuffleVectorInst(Value *Op0, Value *Op1, Constant *Mask,
 
   // Canonicalization: if only one input vector is constant, it shall be the
   // second one.
   if (Op0Const && !Op1Const) {
     std::swap(Op0, Op1);
     ShuffleVectorInst::commuteShuffleMask(Indices, InVecNumElts);
   }
 
+  // A splat of an inserted scalar constant becomes a vector constant:
+  // shuf (inselt ?, C, IndexC), undef, <IndexC, IndexC...> --> <C, C...>
+  // NOTE: We may have commuted above, so analyze the updated Indices, not the
+  //       original mask constant.
+  Constant *C;
+  ConstantInt *IndexC;
+  if (match(Op0, m_InsertElement(m_Value(), m_Constant(C),
+                                 m_ConstantInt(IndexC)))) {
+    // Match a splat shuffle mask of the insert index allowing undef elements.
+    int InsertIndex = IndexC->getZExtValue();
+    if (all_of(Indices, [InsertIndex](int MaskElt) {
+          return MaskElt == InsertIndex || MaskElt == -1;
+        })) {
+      assert(isa<UndefValue>(Op1) && "Expected undef operand 1 for splat");
+
+      // Shuffle mask undefs become undefined constant result elements.
+      SmallVector<Constant *, 16> VecC(MaskNumElts, C);
+      for (unsigned i = 0; i != MaskNumElts; ++i)
+        if (Indices[i] == -1)
+          VecC[i] = UndefValue::get(C->getType());
+      return ConstantVector::get(VecC);
+    }
+  }
+
   // A shuffle of a splat is always the splat itself. Legal if the shuffle's
   // value type is same as the input vectors' type.
   if (auto *OpShuf = dyn_cast<ShuffleVectorInst>(Op0))
     if (isa<UndefValue>(Op1) && RetTy == InVecTy &&
         OpShuf->getMask()->getSplatValue())
       return Op0;
 
   // Don't fold a shuffle with undef mask elements. This may get folded in a

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

;
   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
 }
 
 define <4 x float> @splat_constant(<4 x float> %x) {
 ; CHECK-LABEL: @splat_constant(
 ; CHECK-NEXT:    [[INS3:%.*]] = insertelement <4 x float> [[X:%.*]], float 3.000000e+00, i32 3
-; CHECK-NEXT:    [[SPLAT3:%.*]] = shufflevector <4 x float> [[INS3]], <4 x float> undef, <4 x i32> <i32 3, i32 3, i32 3, i32 3>
-; CHECK-NEXT:    [[R:%.*]] = fadd <4 x float> [[INS3]], [[SPLAT3]]
+; CHECK-NEXT:    [[R:%.*]] = fadd <4 x float> [[INS3]], <float 3.000000e+00, float 3.000000e+00, float 3.000000e+00, float 3.000000e+00>
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %ins3 = insertelement <4 x float> %x, float 3.0, i32 3
   %splat3 = shufflevector <4 x float> %ins3, <4 x float> undef, <4 x i32> <i32 3, i32 3, i32 3, i32 3>
   %r = fadd <4 x float> %ins3, %splat3
   ret <4 x float> %r
 }

llvm/test/Transforms/InstSimplify/shufflevector.ll

 ;
   %tmp1 = shufflevector <2 x float> %x, <2 x float> zeroinitializer, <4 x i32> <i32 2, i32 2, i32 0, i32 1>
   %tmp4 = shufflevector <4 x float> zeroinitializer, <4 x float> %tmp1, <2 x i32> <i32 4, i32 5>
   ret <2 x float> %tmp4
 }
 
 define <5 x i8> @splat_inserted_constant(<4 x i8> %x) {
 ; CHECK-LABEL: @splat_inserted_constant(
-; CHECK-NEXT:    [[INS3:%.*]] = insertelement <4 x i8> [[X:%.*]], i8 42, i64 3
-; CHECK-NEXT:    [[SPLAT5:%.*]] = shufflevector <4 x i8> [[INS3]], <4 x i8> undef, <5 x i32> <i32 3, i32 3, i32 3, i32 3, i32 3>
-; CHECK-NEXT:    ret <5 x i8> [[SPLAT5]]
+; CHECK-NEXT:    ret <5 x i8> <i8 42, i8 42, i8 42, i8 42, i8 42>
 ;
   %ins3 = insertelement <4 x i8> %x, i8 42, i64 3
   %splat5 = shufflevector <4 x i8> %ins3, <4 x i8> undef, <5 x i32> <i32 3, i32 3, i32 3, i32 3, i32 3>
   ret <5 x i8> %splat5
 }
 
 define <4 x float> @splat_inserted_constant_undef_elt(<4 x float> %x) {
 ; CHECK-LABEL: @splat_inserted_constant_undef_elt(
-; CHECK-NEXT:    [[INS1:%.*]] = insertelement <4 x float> [[X:%.*]], float 1.200000e+01, i32 1
-; CHECK-NEXT:    [[SPLAT1:%.*]] = shufflevector <4 x float> [[INS1]], <4 x float> undef, <4 x i32> <i32 1, i32 1, i32 undef, i32 1>
-; CHECK-NEXT:    ret <4 x float> [[SPLAT1]]
+; CHECK-NEXT:    ret <4 x float> <float 1.200000e+01, float 1.200000e+01, float undef, float 1.200000e+01>
 ;
   %ins1 = insertelement <4 x float> %x, float 12.0, i32 1
   %splat1 = shufflevector <4 x float> %ins1, <4 x float> undef, <4 x i32> <i32 1, i32 1, i32 undef, i32 1>
   ret <4 x float> %splat1
 }
 
 define <2 x i8> @splat_inserted_constant_not_canonical(<3 x i8> %x, <3 x i8> %y) {
 ; CHECK-LABEL: @splat_inserted_constant_not_canonical(
-; CHECK-NEXT:    [[INS2:%.*]] = insertelement <3 x i8> [[X:%.*]], i8 23, i7 2
-; CHECK-NEXT:    [[SPLAT2:%.*]] = shufflevector <3 x i8> [[Y:%.*]], <3 x i8> [[INS2]], <2 x i32> <i32 undef, i32 5>
-; CHECK-NEXT:    ret <2 x i8> [[SPLAT2]]
+; CHECK-NEXT:    ret <2 x i8> <i8 undef, i8 23>
 ;
   %ins2 = insertelement <3 x i8> %x, i8 23, i7 2
   %splat2 = shufflevector <3 x i8> %y, <3 x i8> %ins2, <2 x i32> <i32 undef, i32 5>
   ret <2 x i8> %splat2
 }