[InstCombine][SVE] Fix visitExtractElementInst for scalable type.

[huihuiz]

Summary

This patch fix the following issues with visitExtractElementInst:

1. Restrict VectorUtils::findScalarElement to fixed-length vector. For scalable type, the number of elements in shuffle mask is unknown at compile-time.
2. Fix out-of-range calculation for fixed-length vector.
3. Skip scalable type when analysis rely on fixed number of elements.
4. Add unit tests to check functionality of extractelement for scalable type.

[huihuiz]

Current upstream crash at llvm/include/llvm/ADT/APInt.h:1437: void llvm::APInt::setBit(unsigned int): Assertion `BitPosition < BitWidth && "BitPosition out of range"' failed.

take test.ll, run: opt -S -instcombine test.ll

define i32 @extractelement_maybe_out_of_range(<vscale x 4 x i32> %a) {
  %r = extractelement <vscale x 4 x i32> %a, i64 4
  ret i32 %r
}

[huihuiz]

I am also looking at the related issues of shufflevector created by "extractelement, insertelement" pair.
Currently working on InstCombiner::visitInsertElementInst.

[efriedma]

llvm/lib/Analysis/VectorUtils.cpp

305

"EltNo < LHSWidth" doesn't make any sense.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

339–340

I don't think it makes sense to ignore the "scalable" bit here; it's perfectly sensible to extract the third element from a <vscale x 2 x i64>.

[huihuiz]

Thanks Eli for the feedback!
Please let me know if the fixes related to shuffle vector make sense to you.

llvm/lib/Analysis/VectorUtils.cpp

305

Update this fix. I think we should check if EltNo is recorded in shuffle mask before calling getMaskValue.

[efriedma]

llvm/lib/Analysis/VectorUtils.cpp

301

I think I'd prefer to avoid calling getMaskValue() and trying to compute LHSWidth for scalable vectors; it's not clear what it should mean.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

445

Same as above.

[huihuiz]

Address review comments.

[efriedma]

LGTM

[thakis]

Looks like this breaks check-llvm at least on Linux: http://45.33.8.238/linux/17186/step_12.txt

Can you take a look? If it takes a while to fix, please revert while you investigate.

[thakis]

Looks like this breaks check-llvm at least on Linux: http://45.33.8.238/linux/17186/step_12.txt

Can you take a look? If it takes a while to fix, please revert while you investigate.

[huihuiz]

In D78267#2025884, @thakis wrote:

Looks like this breaks check-llvm at least on Linux: http://45.33.8.238/linux/17186/step_12.txt

Can you take a look? If it takes a while to fix, please revert while you investigate.

Already adjusted test check-lines

-; CHECK-NEXT: [[TMP1:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[VEC_E0]], i32 0
+; CHECK-NEXT: [[TMP1:%.*]] = insertelement <vscale x 4 x i32> [[VEC]], i32 [[VEC_E0]], i32 0

Sorry , it wasn't caught by pre-checkin buildbot.

llvm/lib/Analysis/VectorUtils.cpp

if (InsertElementInst *III = dyn_cast<InsertElementInst>(V)) {
     if (EltNo == IIElt)
       return III->getOperand(1);
 
     // Otherwise, the insertelement doesn't modify the value, recurse on its
     // vector input.
     return findScalarElement(III->getOperand(0), EltNo);
   }
 
-  if (ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(V)) {
+  ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(V);
+  // Restrict the following transformation to fixed-length vector.
+  if (SVI && isa<FixedVectorType>(SVI->getType())) {
     unsigned LHSWidth =
-        cast<VectorType>(SVI->getOperand(0)->getType())->getNumElements();
+        cast<FixedVectorType>(SVI->getOperand(0)->getType())->getNumElements();
     int InEl = SVI->getMaskValue(EltNo);
     if (InEl < 0)
       return UndefValue::get(VTy->getElementType());
     if (InEl < (int)LHSWidth)
       return findScalarElement(SVI->getOperand(0), InEl);
     return findScalarElement(SVI->getOperand(1), InEl - LHSWidth);
   }
 
   // Extract a value from a vector add operation with a constant zero.
   // TODO: Use getBinOpIdentity() to generalize this.
   Value *Val; Constant *C;
   if (match(V, m_Add(m_Value(Val), m_Constant(C))))
     if (Constant *Elt = C->getAggregateElement(EltNo))
       if (Elt->isNullValue())
         return findScalarElement(Val, EltNo);
 
   // Otherwise, we don't know.
   return nullptr;

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

Instruction *InstCombiner::visitExtractElementInst(ExtractElementInst &EI) {
   if (Value *V = SimplifyExtractElementInst(SrcVec, Index,
                                             SQ.getWithInstruction(&EI)))
     return replaceInstUsesWith(EI, V);
 
   // If extracting a specified index from the vector, see if we can recursively
   // find a previously computed scalar that was inserted into the vector.
   auto *IndexC = dyn_cast<ConstantInt>(Index);
   if (IndexC) {
-    unsigned NumElts = EI.getVectorOperandType()->getNumElements();
+    ElementCount EC = EI.getVectorOperandType()->getElementCount();
+    unsigned NumElts = EC.Min;
 
     // InstSimplify should handle cases where the index is invalid.
-    if (!IndexC->getValue().ule(NumElts))
+    // For fixed-length vector, it's invalid to extract out-of-range element.
+    if (!EC.Scalable && IndexC->getValue().uge(NumElts))
       return nullptr;
 
     // This instruction only demands the single element from the input vector.
-    if (NumElts != 1) {
+    // Skip for scalable type, the number of elements is unknown at
+    // compile-time.
+    if (!EC.Scalable && NumElts != 1) {
       // If the input vector has a single use, simplify it based on this use
       // property.
       if (SrcVec->hasOneUse()) {
         APInt UndefElts(NumElts, 0);
         APInt DemandedElts(NumElts, 0);
         DemandedElts.setBit(IndexC->getZExtValue());
         if (Value *V =
                 SimplifyDemandedVectorElts(SrcVec, DemandedElts, UndefElts))
if (auto *IE = dyn_cast<InsertElementInst>(I)) {
         return replaceInstUsesWith(EI, IE->getOperand(1));
       // If the inserted and extracted elements are constants, they must not
       // be the same value, extract from the pre-inserted value instead.
       if (isa<Constant>(IE->getOperand(2)) && IndexC)
         return replaceOperand(EI, 0, IE->getOperand(0));
     } else if (auto *SVI = dyn_cast<ShuffleVectorInst>(I)) {
       // If this is extracting an element from a shufflevector, figure out where
       // it came from and extract from the appropriate input element instead.
-      if (auto *Elt = dyn_cast<ConstantInt>(Index)) {
-        int SrcIdx = SVI->getMaskValue(Elt->getZExtValue());
+      // Restrict the following transformation to fixed-length vector.
+      if (isa<FixedVectorType>(SVI->getType()) && isa<ConstantInt>(Index)) {
+        int SrcIdx =
+            SVI->getMaskValue(cast<ConstantInt>(Index)->getZExtValue());
         Value *Src;
-        unsigned LHSWidth =
-            cast<VectorType>(SVI->getOperand(0)->getType())->getNumElements();
+        unsigned LHSWidth = cast<FixedVectorType>(SVI->getOperand(0)->getType())
+                                ->getNumElements();
 
         if (SrcIdx < 0)
           return replaceInstUsesWith(EI, UndefValue::get(EI.getType()));
         if (SrcIdx < (int)LHSWidth)
           Src = SVI->getOperand(0);
         else {
           SrcIdx -= LHSWidth;
           Src = SVI->getOperand(1);
         }
         Type *Int32Ty = Type::getInt32Ty(EI.getContext());
-        return ExtractElementInst::Create(Src,
-                                          ConstantInt::get(Int32Ty,
-                                                           SrcIdx, false));
+        return ExtractElementInst::Create(
+            Src, ConstantInt::get(Int32Ty, SrcIdx, false));
       }
     } else if (auto *CI = dyn_cast<CastInst>(I)) {
       // Canonicalize extractelement(cast) -> cast(extractelement).
       // Bitcasts can change the number of vector elements, and they cost
       // nothing.
       if (CI->hasOneUse() && (CI->getOpcode() != Instruction::BitCast)) {
         Value *EE = Builder.CreateExtractElement(CI->getOperand(0), Index);
         return CastInst::Create(CI->getOpcode(), EE, EI.getType());
       }
     }
   }

llvm/test/Transforms/InstCombine/vscale_extractelement.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i32 @extractelement_in_range(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_in_range(
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i64 1
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %r = extractelement <vscale x 4 x i32> %a, i64 1
+  ret i32 %r
+}
+
+define i32 @extractelement_maybe_out_of_range(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: @extractelement_maybe_out_of_range(
+; 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_bitcast(float %f) {
+; CHECK-LABEL: @extractelement_bitcast(
+; CHECK-NEXT:    [[R:%.*]] = bitcast float [[F:%.*]] to i32
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %vec_float = insertelement <vscale x 4 x float> undef, float %f, i32 0
+  %vec_int = bitcast <vscale x 4 x float> %vec_float to <vscale x 4 x i32>
+  %r = extractelement <vscale x 4 x i32> %vec_int, i32 0
+  ret i32 %r
+}
+
+define i8 @extractelement_bitcast_to_trunc(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_to_trunc(
+; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1
+  %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
+  %r = extractelement <vscale x 8 x i8> %vec_cast, i32 4
+  ret i8 %r
+}
+
+; TODO: Instcombine could remove the insert.
+define i8 @extractelement_bitcast_wrong_insert(<vscale x 2 x i32> %a, i32 %x) {
+; CHECK-LABEL: @extractelement_bitcast_wrong_insert(
+; CHECK-NEXT:    [[VEC:%.*]] = insertelement <vscale x 2 x i32> [[A:%.*]], i32 [[X:%.*]], i32 1
+; CHECK-NEXT:    [[VEC_CAST:%.*]] = bitcast <vscale x 2 x i32> [[VEC]] to <vscale x 8 x i8>
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 8 x i8> [[VEC_CAST]], i32 2
+; CHECK-NEXT:    ret i8 [[R]]
+;
+  %vec = insertelement <vscale x 2 x i32> %a, i32 %x, i32 1 ; <- This insert could be removed.
+  %vec_cast = bitcast <vscale x 2 x i32> %vec to <vscale x 8 x i8>
+  %r = extractelement <vscale x 8 x i8> %vec_cast, i32 2
+  ret i8 %r
+}
+
+; TODO: Instcombine could optimize to return %v.
+define i32 @extractelement_shuffle_in_range(i32 %v) {
+; CHECK-LABEL: @extractelement_shuffle_in_range(
+; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 1
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 1
+  ret i32 %r
+}
+
+define i32 @extractelement_shuffle_maybe_out_of_range(i32 %v) {
+; CHECK-LABEL: @extractelement_shuffle_maybe_out_of_range(
+; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 4
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 4
+  ret i32 %r
+}
+
+define i32 @extractelement_shuffle_invalid_index(i32 %v) {
+; CHECK-LABEL: @extractelement_shuffle_invalid_index(
+; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 -1
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 -1
+  ret i32 %r
+}
+
+
+define i32 @extractelement_shuffle_symbolic_index(i32 %v, i32 %idx) {
+; CHECK-LABEL: @extractelement_shuffle_symbolic_index(
+; CHECK-NEXT:    [[IN:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[V:%.*]], i32 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[IN]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = extractelement <vscale x 4 x i32> [[SPLAT]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 %v, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 %idx
+  ret i32 %r
+}
+
+define <vscale x 4 x i32> @extractelement_insertelement_same_positions(<vscale x 4 x i32> %vec) {
+; CHECK-LABEL: @extractelement_insertelement_same_positions(
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[VEC:%.*]]
+;
+  %vec.e0 = extractelement <vscale x 4 x i32> %vec, i32 0
+  %vec.e1 = extractelement <vscale x 4 x i32> %vec, i32 1
+  %vec.e2 = extractelement <vscale x 4 x i32> %vec, i32 2
+  %vec.e3 = extractelement <vscale x 4 x i32> %vec, i32 3
+  %1 = insertelement <vscale x 4 x i32> %vec, i32 %vec.e0, i32 0
+  %2 = insertelement <vscale x 4 x i32> %1, i32 %vec.e1, i32 1
+  %3 = insertelement <vscale x 4 x i32> %2, i32 %vec.e2, i32 2
+  %4 = insertelement <vscale x 4 x i32> %3, i32 %vec.e3, i32 3
+  ret <vscale x 4 x i32> %4
+}
+
+define <vscale x 4 x i32> @extractelement_insertelement_diff_positions(<vscale x 4 x i32> %vec) {
+; CHECK-LABEL: @extractelement_insertelement_diff_positions(
+; CHECK-NEXT:    [[VEC_E0:%.*]] = extractelement <vscale x 4 x i32> [[VEC:%.*]], i32 4
+; CHECK-NEXT:    [[VEC_E1:%.*]] = extractelement <vscale x 4 x i32> [[VEC]], i32 5
+; CHECK-NEXT:    [[VEC_E2:%.*]] = extractelement <vscale x 4 x i32> [[VEC]], i32 6
+; CHECK-NEXT:    [[VEC_E3:%.*]] = extractelement <vscale x 4 x i32> [[VEC]], i32 7
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <vscale x 4 x i32> undef, i32 [[VEC_E0]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <vscale x 4 x i32> [[TMP1]], i32 [[VEC_E1]], i32 1
+; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <vscale x 4 x i32> [[TMP2]], i32 [[VEC_E2]], i32 2
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <vscale x 4 x i32> [[TMP3]], i32 [[VEC_E3]], i32 3
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[TMP4]]
+;
+  %vec.e0 = extractelement <vscale x 4 x i32> %vec, i32 4
+  %vec.e1 = extractelement <vscale x 4 x i32> %vec, i32 5
+  %vec.e2 = extractelement <vscale x 4 x i32> %vec, i32 6
+  %vec.e3 = extractelement <vscale x 4 x i32> %vec, i32 7
+  %1 = insertelement <vscale x 4 x i32> %vec, i32 %vec.e0, i32 0
+  %2 = insertelement <vscale x 4 x i32> %1, i32 %vec.e1, i32 1
+  %3 = insertelement <vscale x 4 x i32> %2, i32 %vec.e2, i32 2
+  %4 = insertelement <vscale x 4 x i32> %3, i32 %vec.e3, i32 3
+  ret <vscale x 4 x i32> %4
+}