Constfold insertelement to undef when index is out-of-bounds

[chfast]

Summary

This patch adds constant folding of insertelement instruction to undef value when index operand is constant and is not less than vector size or is undef.

InstCombine does not support this case, but I'm happy to add it there also if this change is accepted.

[majnemer]

lib/IR/ConstantFold.cpp

828

It might be nicer to do something like:

uint64_t Idx = CIdx->getZExtValue();

After the above comparison, we know that getZExtValue will succeed. This lets us simplify the loop below a little (i.e. Idx == i vs CIdx->equalsInt(i))

829

No need to initialize the SmallVector. Please use Result.reserve(NumElts); instead.

832–833

This will need to switch to using push_back once you use reserve.

[chfast]

Switched back to push_back driven implementation

[majnemer]

LGTM

[chfast]

Thanks for the review.

What do you think about adding "the same" to InstCombine?

[majnemer]

InstCombine would work but there is a better place: InstSimplify. It is
preferable to stick things in InstSimplify when a new instruction doesn't
need to be created.

lib/IR/ConstantFold.cpp

if (ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx)) {
     return Val->getAggregateElement(CIdx->getZExtValue());
   }
   return nullptr;
 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
                                                      Constant *Elt,
                                                      Constant *Idx) {
+  if (isa<UndefValue>(Idx))
+    return UndefValue::get(Val->getType());
+
   ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx) return nullptr;
-  const APInt &IdxVal = CIdx->getValue();
-  
+
+  unsigned NumElts = Val->getType()->getVectorNumElements();
+  if (CIdx->uge(NumElts))
+    return UndefValue::get(Val->getType());
   SmallVector<Constant*, 16> Result;
-  Type *Ty = IntegerType::get(Val->getContext(), 32);
-  for (unsigned i = 0, e = Val->getType()->getVectorNumElements(); i != e; ++i){
+  Result.reserve(NumElts);
+  auto *Ty = Type::getInt32Ty(Val->getContext());
+  uint64_t IdxVal = CIdx->getZExtValue();
+  for (unsigned i = 0; i != NumElts; ++i) {    
     if (i == IdxVal) {
       Result.push_back(Elt);
       continue;
     }
     
-    Constant *C =
-      ConstantExpr::getExtractElement(Val, ConstantInt::get(Ty, i));
+    Constant *C = ConstantExpr::getExtractElement(Val, ConstantInt::get(Ty, i));
     Result.push_back(C);
   }
-  
   return ConstantVector::get(Result);
 }
 
 Constant *llvm::ConstantFoldShuffleVectorInstruction(Constant *V1,
                                                      Constant *V2,
                                                      Constant *Mask) {
   unsigned MaskNumElts = Mask->getType()->getVectorNumElements();
   Type *EltTy = V1->getType()->getVectorElementType();
 
   // Undefined shuffle mask -> undefined value.
   if (isa<UndefValue>(Mask))
     return UndefValue::get(VectorType::get(EltTy, MaskNumElts));
 
   // Don't break the bitcode reader hack.
   if (isa<ConstantExpr>(Mask)) return nullptr;

test/Transforms/ConstProp/InsertElement.ll

 ; RUN: opt < %s -constprop -S | FileCheck %s
 
+; CHECK-LABEL: @test1
 define i32 @test1() {
   %A = bitcast i32 2139171423 to float
   %B = insertelement <1 x float> undef, float %A, i32 0
   %C = extractelement <1 x float> %B, i32 0
   %D = bitcast float %C to i32
   ret i32 %D
-; CHECK: @test1
 ; CHECK: ret i32 2139171423
 }
 
+; CHECK-LABEL: @insertelement
+define <4 x i64> @insertelement() {
+  %vec1 = insertelement <4 x i64> undef, i64 -1, i32 0
+  %vec2 = insertelement <4 x i64> %vec1, i64 -2, i32 1
+  %vec3 = insertelement <4 x i64> %vec2, i64 -3, i32 2
+  %vec4 = insertelement <4 x i64> %vec3, i64 -4, i32 3
+  ; CHECK: ret <4 x i64> <i64 -1, i64 -2, i64 -3, i64 -4>
+  ret <4 x i64> %vec4
+}
+
+; CHECK-LABEL: @insertelement_undef
+define <4 x i64> @insertelement_undef() {
+  %vec1 = insertelement <4 x i64> undef, i64 -1, i32 0
+  %vec2 = insertelement <4 x i64> %vec1, i64 -2, i32 1
+  %vec3 = insertelement <4 x i64> %vec2, i64 -3, i32 2
+  %vec4 = insertelement <4 x i64> %vec3, i64 -4, i32 3
+  %vec5 = insertelement <4 x i64> %vec3, i64 -5, i32 4
+  ; CHECK: ret <4 x i64> undef
+  ret <4 x i64> %vec5
+}

unittests/IR/ConstantsTest.cpp

TEST(ConstantsTest, AsInstructionsTest) {
   Constant *P4 = ConstantExpr::getPtrToInt(Global2, Int32Ty);
   Constant *P5 = ConstantExpr::getUIToFP(P4, FloatTy);
   Constant *P6 = ConstantExpr::getBitCast(P4, VectorType::get(Int16Ty, 2));
 
   Constant *One = ConstantInt::get(Int32Ty, 1);
   Constant *Two = ConstantInt::get(Int64Ty, 2);
   Constant *Big = ConstantInt::get(getGlobalContext(),
                                    APInt{256, uint64_t(-1), true});
-  Constant *Undef = UndefValue::get(Int64Ty);
+  Constant *Elt = ConstantInt::get(Int16Ty, 2015);
+  Constant *Undef16  = UndefValue::get(Int16Ty);
+  Constant *Undef64  = UndefValue::get(Int64Ty);
+  Constant *UndefV16 = UndefValue::get(P6->getType());
 
   #define P0STR "ptrtoint (i32** @dummy to i32)"
   #define P1STR "uitofp (i32 ptrtoint (i32** @dummy to i32) to float)"
   #define P2STR "uitofp (i32 ptrtoint (i32** @dummy to i32) to double)"
   #define P3STR "ptrtoint (i32** @dummy to i1)"
   #define P4STR "ptrtoint (i32** @dummy2 to i32)"
   #define P5STR "uitofp (i32 ptrtoint (i32** @dummy2 to i32) to float)"
   #define P6STR "bitcast (i32 ptrtoint (i32** @dummy2 to i32) to <2 x i16>)"
TEST(ConstantsTest, AsInstructionsTest) {
   //      "getelementptr i32*, i32** @dummy, i32 1");
   CHECK(ConstantExpr::getInBoundsGetElementPtr(PointerType::getUnqual(Int32Ty),
                                                Global, V),
         "getelementptr inbounds i32*, i32** @dummy, i32 1");
 
   CHECK(ConstantExpr::getExtractElement(P6, One), "extractelement <2 x i16> "
         P6STR ", i32 1");
 
-  EXPECT_TRUE(isa<UndefValue>(ConstantExpr::getExtractElement(P6, Two)));
-  EXPECT_TRUE(isa<UndefValue>(ConstantExpr::getExtractElement(P6, Big)));
-  EXPECT_TRUE(isa<UndefValue>(ConstantExpr::getExtractElement(P6, Undef)));
+  EXPECT_EQ(Undef16, ConstantExpr::getExtractElement(P6, Two));
+  EXPECT_EQ(Undef16, ConstantExpr::getExtractElement(P6, Big));
+  EXPECT_EQ(Undef16, ConstantExpr::getExtractElement(P6, Undef64));
+
+  EXPECT_EQ(Elt, ConstantExpr::getExtractElement(
+                 ConstantExpr::getInsertElement(P6, Elt, One), One));
+  EXPECT_EQ(UndefV16, ConstantExpr::getInsertElement(P6, Elt, Two));
+  EXPECT_EQ(UndefV16, ConstantExpr::getInsertElement(P6, Elt, Big));
+  EXPECT_EQ(UndefV16, ConstantExpr::getInsertElement(P6, Elt, Undef64));
 }
 
 #ifdef GTEST_HAS_DEATH_TEST
 #ifndef NDEBUG
 TEST(ConstantsTest, ReplaceWithConstantTest) {
   std::unique_ptr<Module> M(new Module("MyModule", getGlobalContext()));
 
   Type *Int32Ty = Type::getInt32Ty(getGlobalContext());