[IR] Allow constant folding (insertelement <vscale x 2 x i32> zeroinitializer, i32 0, i32 i32 0.

[craig.topper]

Summary

Most of insertelement constant folding is blocked if the vector type
is scalable. I believe we can make an exception for inserting null
into an all zeros vector.

[nikic]

llvm/lib/IR/ConstantFold.cpp

696

This could go before the CIdx check. Also, this holds fold any splat Val (e.g. UndefValue or PoisonValue), not just zero.

[craig.topper]

Move earlier. Add a TODO for undef/poison.

[craig.topper]

Ping

[nikic]

LG, but please add a dedicated test for this (in test/Transforms/InstSimplify/ConstantFold or something).

[reames]

llvm/lib/IR/ConstantFold.cpp

687

I believe this should hold for any insertion of a constant into a splat of the same value. Poison/Undef/Zero are all just special cases of that.

llvm/lib/IR/ConstantFold.cpp

 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
                                                      Constant *Elt,
                                                      Constant *Idx) {
   if (isa<UndefValue>(Idx))
     return PoisonValue::get(Val->getType());
 
+  // Inserting null into all zeros is still all zeros.
+  // TODO: This is true for undef and poison splats too.
+  if (isa<ConstantAggregateZero>(Val) && Elt->isNullValue())
+    return Val;
+
   ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx) return nullptr;
 
   // Do not iterate on scalable vector. The num of elements is unknown at
   // compile-time.
   if (isa<ScalableVectorType>(Val->getType()))
     return nullptr;
 
   auto *ValTy = cast<FixedVectorType>(Val->getType());
 
   unsigned NumElts = ValTy->getNumElements();
   if (CIdx->uge(NumElts))
     return PoisonValue::get(Val->getType());
 

llvm/test/Transforms/InstSimplify/ConstProp/InsertElement.ll

 
 define i64 @extract_undef_index_from_nonzero_vec() {
 ; CHECK-LABEL: @extract_undef_index_from_nonzero_vec(
 ; CHECK-NEXT:    ret i64 poison
 ;
   %E = extractelement <2 x i64> <i64 -1, i64 -1>, i64 undef
   ret i64 %E
 }
+
+define <vscale x 2 x i32> @insertelement_scalable_null() {
+; CHECK-LABEL: @insertelement_scalable_null(
+; CHECK-NEXT:    ret <vscale x 2 x i32> zeroinitializer
+;
+  %vec = insertelement <vscale x 2 x i32> zeroinitializer, i32 0, i32 0
+  ret <vscale x 2 x i32> %vec
+}
+
+define <vscale x 2 x float> @insertelement_scalable_null_fp() {
+; CHECK-LABEL: @insertelement_scalable_null_fp(
+; CHECK-NEXT:    ret <vscale x 2 x float> zeroinitializer
+;
+  %vec = insertelement <vscale x 2 x float> zeroinitializer, float 0.0, i32 1
+  ret <vscale x 2 x float> %vec
+}

llvm/test/Transforms/LoopVectorize/AArch64/sve-widen-phi.ll

 ; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i64 [[SMAX]], [[TMP3]]
 ; CHECK-NEXT:    [[N_VEC:%.*]] = sub nsw i64 [[SMAX]], [[N_MOD_VF]]
 ; CHECK-NEXT:    [[IND_END:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 [[N_VEC]]
 ; CHECK-NEXT:    [[IND_END3:%.*]] = getelementptr i32*, i32** [[B:%.*]], i64 [[N_VEC]]
 ; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[POINTER_PHI:%.*]] = phi i32* [ [[A]], [[VECTOR_PH]] ], [ [[PTR_IND:%.*]], [[VECTOR_BODY]] ]
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <vscale x 2 x i32> [ insertelement (<vscale x 2 x i32> zeroinitializer, i32 0, i32 0), [[VECTOR_PH]] ], [ [[TMP9:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <vscale x 2 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP9:%.*]], [[VECTOR_BODY]] ]
 ; CHECK-NEXT:    [[TMP4:%.*]] = call i64 @llvm.vscale.i64()
 ; CHECK-NEXT:    [[TMP5:%.*]] = shl nuw nsw i64 [[TMP4]], 1
 ; CHECK-NEXT:    [[TMP6:%.*]] = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
 ; CHECK-NEXT:    [[TMP7:%.*]] = getelementptr i32, i32* [[POINTER_PHI]], <vscale x 2 x i64> [[TMP6]]
 ; CHECK-NEXT:    [[NEXT_GEP:%.*]] = getelementptr i32*, i32** [[B]], i64 [[INDEX]]
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast <vscale x 2 x i32*> [[TMP7]] to <vscale x 2 x <vscale x 2 x i32>*>
 ; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <vscale x 2 x <vscale x 2 x i32>*> [[BC]], i64 0
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <vscale x 2 x i32>, <vscale x 2 x i32>* [[TMP8]], align 8