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

[huihuiz]

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.

[huihuiz]

Current upstream is doing wrong fold for scalable vectors.

for test case: insertelement_idx_maybe_out_of_bound

define <vscale x 4 x i32> @insertelement_idx_maybe_out_of_bound(<vscale x 4 x i32> %a) {
  %r = insertelement <vscale x 4 x i32> %a, i32 5, i64 4
  ret <vscale x 4 x i32> %r
}

run: opt -instsimplify -S t.ll -o -

we end up with

define <vscale x 4 x i32> @insertelement_idx_maybe_out_of_bound(<vscale x 4 x i32> %a) {
  ret <vscale x 4 x i32> undef
}

Index 4 could be out-of-bound for scalable vector, however we don't know this at compile-time.

Also take test case insert_extract_element_same_vec_idx_2

define i32 @insert_extract_element_same_vec_idx_2(<vscale x 4 x i32> %a) {
  %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
}

run: opt -instsimplify -S t.ll -o -

we got

define i32 @insert_extract_element_same_vec_idx_2(<vscale x 4 x i32> %a) {
  ret i32 undef
}

That's the wrong fold in findScalarElement

we should be getting

define i32 @insert_extract_element_same_vec_idx_2(<vscale x 4 x i32> %a) {
  ret i32 1
}

[efriedma]

LGTM

llvm/lib/Analysis/InstructionSimplify.cpp

Value *llvm::SimplifyInsertElementInst(Value *Vec, Value *Val, Value *Idx,
                                        const SimplifyQuery &Q) {
   // Try to constant fold.
   auto *VecC = dyn_cast<Constant>(Vec);
   auto *ValC = dyn_cast<Constant>(Val);
   auto *IdxC = dyn_cast<Constant>(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());
   }
 
   // If index is undef, it might be out of bounds (see above case)
   if (isa<UndefValue>(Idx))
     return UndefValue::get(Vec->getType());
 
   // Inserting an undef scalar? Assume it is the same value as the existing
if (auto *CVec = dyn_cast<Constant>(Vec)) {
 
     if (isa<UndefValue>(Vec))
       return UndefValue::get(Vec->getType()->getVectorElementType());
   }
 
   // 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;
   }
 
   // An undef extract index can be arbitrarily chosen to be an out-of-range
   // index value, which would result in the instruction being undef.
   if (isa<UndefValue>(Idx))

llvm/lib/Analysis/VectorUtils.cpp

 }
 
 /// Given a vector and an element number, see if the scalar value is
 /// already around as a register, for example if it were inserted then extracted
 /// from the vector.
 Value *llvm::findScalarElement(Value *V, unsigned EltNo) {
   assert(V->getType()->isVectorTy() && "Not looking at a vector?");
   VectorType *VTy = cast<VectorType>(V->getType());
+  // For fixed-length vector, return undef for out of range access.
+  if (!V->getType()->getVectorIsScalable()) {
     unsigned Width = VTy->getNumElements();
-  if (EltNo >= Width)  // Out of range access.
+    if (EltNo >= Width)
       return UndefValue::get(VTy->getElementType());
+  }
 
   if (Constant *C = dyn_cast<Constant>(V))
     return C->getAggregateElement(EltNo);
 
   if (InsertElementInst *III = dyn_cast<InsertElementInst>(V)) {
     // If this is an insert to a variable element, we don't know what it is.
     if (!isa<ConstantInt>(III->getOperand(2)))
       return nullptr;

llvm/test/Transforms/InstSimplify/vscale.ll

+; 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
+}
+