[InstSimplify][SVE] Fix SimplifyInsert/ExtractElementInst for scalable vector.

Summary:
For scalable vector, index out-of-bound can not be determined at compile-time.
The same apply for VectorUtil findScalarElement().

Add test cases to check the functionality of SimplifyInsert/ExtractElementInst for scalable vector.

Reviewers: sdesmalen, efriedma, spatel, apazos

Reviewed By: efriedma

Subscribers: cameron.mcinally, tschuett, hiraditya, rkruppe, psnobl, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75782

3 files changed, 108 insertions, 8 deletions

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 47f0340..ef3ae1f 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -4220,10 +4220,10 @@ Value *llvm::SimplifyInsertElementInst(Value *Vec, Value *Val, Value *Idx,
   if (VecC && ValC && IdxC)
     return ConstantFoldInsertElementInstruction(VecC, ValC, IdxC);
 
-  // Fold into undef if index is out of bounds.
+  // For fixed-length vector, 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))
+    if (!Vec->getType()->getVectorIsScalable() &&
+        CI->uge(Vec->getType()->getVectorNumElements()))
       return UndefValue::get(Vec->getType());
   }
 
@@ -4294,8 +4294,9 @@ 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().uge(Vec->getType()->getVectorNumElements()))
-      // definitely out of bounds, thus undefined result
+    // For fixed-length vector, fold into undef if index is out of bounds.
+    if (!Vec->getType()->getVectorIsScalable() &&
+        IdxC->getValue().uge(Vec->getType()->getVectorNumElements()))
       return UndefValue::get(Vec->getType()->getVectorElementType());
     if (Value *Elt = findScalarElement(Vec, IdxC->getZExtValue()))
       return Elt;
diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index 2d0c85c..532518a 100644
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -262,9 +262,12 @@ Value *llvm::getStrideFromPointer(Value *Ptr, ScalarEvolution *SE, Loop *Lp) {
 Value *llvm::findScalarElement(Value *V, unsigned EltNo) {
   assert(V->getType()->isVectorTy() && "Not looking at a vector?");
   VectorType *VTy = cast<VectorType>(V->getType());
-  unsigned Width = VTy->getNumElements();
-  if (EltNo >= Width)  // Out of range access.
-    return UndefValue::get(VTy->getElementType());
+  // For fixed-length vector, return undef for out of range access.
+  if (!V->getType()->getVectorIsScalable()) {
+    unsigned Width = VTy->getNumElements();
+    if (EltNo >= Width)
+      return UndefValue::get(VTy->getElementType());
+  }
 
   if (Constant *C = dyn_cast<Constant>(V))
     return C->getAggregateElement(EltNo);
diff --git a/llvm/test/Transforms/InstSimplify/vscale.ll b/llvm/test/Transforms/InstSimplify/vscale.ll
new file mode 100644
index 0000000..5a6971c
--- /dev/null
+++ b/llvm/test/Transforms/InstSimplify/vscale.ll
@@ -0,0 +1,96 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instsimplify -S -verify | FileCheck %s
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Vector Operations
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+; insertelement
+
+define <vscale x 4 x i32> @insertelement_idx_undef(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insertelement_idx_undef(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = insertelement <vscale x 4 x i32> %a, i32 5, i64 undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @insertelement_value_undef(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insertelement_value_undef(
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[A:%.*]]
+;
+  %r = insertelement <vscale x 4 x i32> %a, i32 undef, i64 0
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @insertelement_idx_maybe_out_of_bound(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insertelement_idx_maybe_out_of_bound(
+; CHECK-NEXT:    [[R:%.*]] = insertelement <vscale x 4 x i32> [[A:%.*]], i32 5, i64 4
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[R]]
+;
+  %r = insertelement <vscale x 4 x i32> %a, i32 5, i64 4
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @insertelement_idx_large_bound(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insertelement_idx_large_bound(
+; CHECK-NEXT:    [[R:%.*]] = insertelement <vscale x 4 x i32> [[A:%.*]], i32 5, i64 12345
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[R]]
+;
+  %r = insertelement <vscale x 4 x i32> %a, i32 5, i64 12345
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @insert_extract_element_same_vec_idx_1(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insert_extract_element_same_vec_idx_1(
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[A:%.*]]
+;
+  %v = extractelement <vscale x 4 x i32> %a, i64 1
+  %r = insertelement <vscale x 4 x i32> %a, i32 %v, i64 1
+  ret <vscale x 4 x i32> %r
+}
+
+; extractelement
+
+define i32 @extractelement_idx_undef(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_idx_undef(
+; CHECK-NEXT:    ret i32 undef
+;
+  %r = extractelement <vscale x 4 x i32> %a, i64 undef
+  ret i32 %r
+}
+
+define i32 @extractelement_vec_undef(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_vec_undef(
+; CHECK-NEXT:    ret i32 undef
+;
+  %r = extractelement <vscale x 4 x i32> undef, i64 1
+  ret i32 %r
+}
+
+define i32 @extractelement_idx_maybe_out_of_bound(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_idx_maybe_out_of_bound(
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i64 4
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %r = extractelement <vscale x 4 x i32> %a, i64 4
+  ret i32 %r
+}
+define i32 @extractelement_idx_large_bound(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_idx_large_bound(
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i64 12345
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %r = extractelement <vscale x 4 x i32> %a, i64 12345
+  ret i32 %r
+}
+
+define i32 @insert_extract_element_same_vec_idx_2(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @insert_extract_element_same_vec_idx_2(
+; CHECK-NEXT:    ret i32 1
+;
+  %v = insertelement <vscale x 4 x i32> undef, i32 1, i64 4
+  %r = extractelement <vscale x 4 x i32> %v, i64 4
+  ret i32 %r
+}
+