[instsimplify] consistently handle undef and out of bound indices for insertelement and extractelement

In one case, we were handling out of bounds, but not undef indices.  In the other, we were handling undef (with the comment making the analogy to out of bounds), but not out of bounds.  Be consistent and treat both undef and constant out of bounds indices as producing undefined results.

As a side effect, this also protects instcombine from having to handle large constant indices as we always simplify first.

6 files changed, 57 insertions, 16 deletions

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 93fb1143e50..c94429f5521 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -3838,12 +3838,13 @@ Value *llvm::SimplifyInsertElementInst(Value *Vec, Value *Val, Value *Idx,
   // Fold into undef if index is out of bounds.
   if (auto *CI = dyn_cast<ConstantInt>(Idx)) {
     uint64_t NumElements = cast<VectorType>(Vec->getType())->getNumElements();
-
     if (CI->uge(NumElements))
       return UndefValue::get(Vec->getType());
   }
 
-  // TODO: We should also fold if index is iteslf an undef.
+  // If index is undef, it might be out of bounds (see above case)
+  if (isa<UndefValue>(Idx))
+    return UndefValue::get(Vec->getType());
 
   return nullptr;
 }
@@ -3896,10 +3897,13 @@ static Value *SimplifyExtractElementInst(Value *Vec, Value *Idx, const SimplifyQ
 
   // If extracting a specified index from the vector, see if we can recursively
   // find a previously computed scalar that was inserted into the vector.
-  if (auto *IdxC = dyn_cast<ConstantInt>(Idx))
-    if (IdxC->getValue().ule(Vec->getType()->getVectorNumElements()))
-      if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
-        return Elt;
+  if (auto *IdxC = dyn_cast<ConstantInt>(Idx)) {
+    if (IdxC->getValue().uge(Vec->getType()->getVectorNumElements()))
+      // definitely out of bounds, thus undefined result
+      return UndefValue::get(Vec->getType()->getVectorElementType());
+    if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
+      return Elt;
+  }
 
   // An undef extract index can be arbitrarily chosen to be an out-of-range
   // index value, which would result in the instruction being undef.
diff --git a/llvm/test/Transforms/InstCombine/extractelement.ll b/llvm/test/Transforms/InstCombine/extractelement.ll
index 242f0623c2c..f4043335c4e 100644
--- a/llvm/test/Transforms/InstCombine/extractelement.ll
+++ b/llvm/test/Transforms/InstCombine/extractelement.ll
@@ -3,8 +3,7 @@
 
 define i32 @extractelement_out_of_range(<2 x i32> %x) {
 ; CHECK-LABEL: @extractelement_out_of_range(
-; CHECK-NEXT:    [[E1:%.*]] = extractelement <2 x i32> [[X:%.*]], i8 16
-; CHECK-NEXT:    ret i32 [[E1]]
+; CHECK-NEXT:    ret i32 undef
 ;
   %E1 = extractelement <2 x i32> %x, i8 16
   ret i32 %E1
diff --git a/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll b/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
index 9d59efbad73..318df6cf76c 100644
--- a/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
+++ b/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
@@ -191,11 +191,11 @@ define <4 x i32> @inselt_shuf_no_demand_multiuse(i32 %a0, i32 %a1, <4 x i32> %b)
 
 define <4 x float> @inselt_shuf_no_demand_bogus_insert_index_in_chain(float %a1, float %a2, float %a3, i32 %variable_index) {
 ; CHECK-LABEL: @inselt_shuf_no_demand_bogus_insert_index_in_chain(
-; CHECK-NEXT:    [[OUT1:%.*]] = insertelement <4 x float> undef, float %a1, i32 1
-; CHECK-NEXT:    ret <4 x float> [[OUT1]]
+; CHECK-NEXT:    [[OUT12:%.*]] = insertelement <4 x float> undef, float [[A2:%.*]], i32 [[VARIABLE_INDEX:%.*]]
+; CHECK-NEXT:    ret <4 x float> [[OUT12]]
 ;
   %out1 = insertelement <4 x float> undef, float %a1, i32 1
-  %out12 = insertelement <4 x float> %out1, float %a2, i32 undef ; something unexpected
+  %out12 = insertelement <4 x float> %out1, float %a2, i32 %variable_index ; something unexpected
   %out123 = insertelement <4 x float> %out12, float %a3, i32 3
   %shuffle = shufflevector <4 x float> %out123, <4 x float> undef, <4 x i32> <i32 0, i32 undef, i32 undef, i32 undef>
   ret <4 x float> %shuffle
diff --git a/llvm/test/Transforms/InstCombine/vector_insertelt_shuffle.ll b/llvm/test/Transforms/InstCombine/vector_insertelt_shuffle.ll
index 41c6370e48e..e5da6086319 100644
--- a/llvm/test/Transforms/InstCombine/vector_insertelt_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vector_insertelt_shuffle.ll
@@ -66,9 +66,7 @@ define <4 x float> @bazzz(<4 x float> %x) {
 
 define <4 x float> @bazzzz(<4 x float> %x) {
 ; CHECK-LABEL: @bazzzz(
-; CHECK-NEXT:    [[INS1:%.*]] = insertelement <4 x float> %x, float 1.000000e+00, i32 undef
-; CHECK-NEXT:    [[INS2:%.*]] = insertelement <4 x float> %x, float 2.000000e+00, i32 2
-; CHECK-NEXT:    ret <4 x float> [[INS2]]
+; CHECK-NEXT:   ret <4 x float> <float undef, float undef, float 2.000000e+00, float undef>
 ;
   %ins1 = insertelement<4 x float> %x, float 1.0, i32 undef
   %ins2 = insertelement<4 x float> %ins1, float 2.0, i32 2
diff --git a/llvm/test/Transforms/InstSimplify/extract-element.ll b/llvm/test/Transforms/InstSimplify/extract-element.ll
index 8ee75a603cd..05147891312 100644
--- a/llvm/test/Transforms/InstSimplify/extract-element.ll
+++ b/llvm/test/Transforms/InstSimplify/extract-element.ll
@@ -5,9 +5,43 @@
 
 define i129 @vec_extract_negidx(<3 x i129> %a) {
 ; CHECK-LABEL: @vec_extract_negidx(
-; CHECK-NEXT:    [[E1:%.*]] = extractelement <3 x i129> [[A:%.*]], i129 -1
-; CHECK-NEXT:    ret i129 [[E1]]
+; CHECK-NEXT:    ret i129 undef
 ;
   %E1 = extractelement <3 x i129> %a, i129 -1
   ret i129 %E1
 }
+
+define i129 @vec_extract_out_of_bounds(<3 x i129> %a) {
+; CHECK-LABEL: @vec_extract_out_of_bounds(
+; CHECK-NEXT:    ret i129 undef
+;
+  %E1 = extractelement <3 x i129> %a, i129 3
+  ret i129 %E1
+}
+
+define i129 @vec_extract_out_of_bounds2(<3 x i129> %a) {
+; CHECK-LABEL: @vec_extract_out_of_bounds2(
+; CHECK-NEXT:    ret i129 undef
+;
+  %E1 = extractelement <3 x i129> %a, i129 999999999999999
+  ret i129 %E1
+}
+
+
+define i129 @vec_extract_undef_index(<3 x i129> %a) {
+; CHECK-LABEL: @vec_extract_undef_index(
+; CHECK-NEXT:    ret i129 undef
+;
+  %E1 = extractelement <3 x i129> %a, i129 undef
+  ret i129 %E1
+}
+
+
+define i129 @vec_extract_in_bounds(<3 x i129> %a) {
+; CHECK-LABEL: @vec_extract_in_bounds(
+; CHECK-NEXT:    %E1 = extractelement <3 x i129> %a, i129 2
+; CHECK-NEXT:     ret i129 %E1
+;
+  %E1 = extractelement <3 x i129> %a, i129 2
+  ret i129 %E1
+}
diff --git a/llvm/test/Transforms/InstSimplify/insertelement.ll b/llvm/test/Transforms/InstSimplify/insertelement.ll
index 3acd921cbad..3524f2145ac 100644
--- a/llvm/test/Transforms/InstSimplify/insertelement.ll
+++ b/llvm/test/Transforms/InstSimplify/insertelement.ll
@@ -23,3 +23,9 @@ define <4 x i32> @test4(<4 x i32> %A) {
   ; CHECK: ret <4 x i32> undef
   ret <4 x i32> %I
 }
+
+define <4 x i32> @test5(<4 x i32> %A) {
+  %I = insertelement <4 x i32> %A, i32 5, i64 undef
+  ; CHECK: ret <4 x i32> undef
+  ret <4 x i32> %I
+}