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

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.

Test Plan: Unittests and regression tests for ConstProp pass.

Reviewers: majnemer

Reviewed By: majnemer

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D9287

llvm-svn: 235854

3 files changed, 48 insertions, 12 deletions

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index 03d7c5e7b95..2a524937391 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -800,23 +800,30 @@ Constant *llvm::ConstantFoldExtractElementInstruction(Constant *Val,
 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);
 }
 
diff --git a/llvm/test/Transforms/ConstProp/InsertElement.ll b/llvm/test/Transforms/ConstProp/InsertElement.ll
index d249c2eb1c5..011ad3f57f4 100644
--- a/llvm/test/Transforms/ConstProp/InsertElement.ll
+++ b/llvm/test/Transforms/ConstProp/InsertElement.ll
@@ -1,12 +1,32 @@
 ; 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
+}
diff --git a/llvm/unittests/IR/ConstantsTest.cpp b/llvm/unittests/IR/ConstantsTest.cpp
index d02217dc355..7741b448fa8 100644
--- a/llvm/unittests/IR/ConstantsTest.cpp
+++ b/llvm/unittests/IR/ConstantsTest.cpp
@@ -190,7 +190,10 @@ TEST(ConstantsTest, AsInstructionsTest) {
   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)"
@@ -260,9 +263,15 @@ TEST(ConstantsTest, AsInstructionsTest) {
   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