[SLP] Simplify indices processing for insertelements

[anton-afanasyev]

Summary

Get rid of non-constant and undef indices of insertelements
at buildTree() stage. Fix bugs.

[RKSimon]

Something to take it account with this and D119536 - we're going to need to cherrypick these fixes for PR52275 as they affect 14.x as well - will they be able to apply cleanly to the 14.x branch?

[ABataev]

What kind of bugs do you see with these test cases and current version of the function? I would try to keep it, better to generate UndefMaskElem if possible, rather than cut of the vectorization early.

[anton-afanasyev]

In D119623#3316827, @ABataev wrote:

What kind of bugs do you see with these test cases and current version of the function? I would try to keep it, better to generate UndefMaskElem if possible, rather than cut of the vectorization early.

All these test cases are crashing without patch. There are several issues with this code. First of all, see lines 4028 and 4053 -- we use already dereferenced Optional<int> Idx = *getInsertIndex(V); (it should be without *). Also MinIdx is computed but not used at all.

But the main problem is that UndefMaskElem generally makes no sense for index: the function getInsertIndex() returns _index_ rather than element, if it is undef/poison/out of range then whole vector is poison. Also this return value UndefMaskElem isn't used now anywhere, if somehow (how?) it could be useful in future, it can be added.

[anton-afanasyev]

In D119623#3316685, @RKSimon wrote:

Something to take it account with this and D119536 - we're going to need to cherrypick these fixes for PR52275 as they affect 14.x as well - will they be able to apply cleanly to the 14.x branch?

No, they doesn't apply cleanly right now, but I'm to backport them through new workflow (https://discourse.llvm.org/t/rfc-new-automated-release-workflow-part-2/59981)

[ABataev]

In D119623#3317456, @anton-afanasyev wrote:

In D119623#3316827, @ABataev wrote:

What kind of bugs do you see with these test cases and current version of the function? I would try to keep it, better to generate UndefMaskElem if possible, rather than cut of the vectorization early.

All these test cases are crashing without patch. There are several issues with this code. First of all, see lines 4028 and 4053 -- we use already dereferenced Optional<int> Idx = *getInsertIndex(V); (it should be without *). Also MinIdx is computed but not used at all.

But the main problem is that UndefMaskElem generally makes no sense for index: the function getInsertIndex() returns _index_ rather than element, if it is undef/poison/out of range then whole vector is poison. Also this return value UndefMaskElem isn't used now anywhere, if somehow (how?) it could be useful in future, it can be added.

Ok, got it. Could you also add a check to the buildvector detection function for such kind of insertelement (+ possibly insertvalue) instructions? They should not be treated as a part of buildvector, need to treat them as an initial/terminating instruction and exclude from possibly vectorizable buildvector instructions.

[anton-afanasyev]

Ok, got it. Could you also add a check to the buildvector detection function for such kind of insertelement (+ possibly insertvalue) instructions? They should not be treated as a part of buildvector, need to treat them as an initial/terminating instruction and exclude from possibly vectorizable buildvector instructions.

Do you mean findBuildAggregate()? This check is already there, see line 9603.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9603

@ABataev Do you mean this check?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9603

Yes, but I thought about returning true here.

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9603

Hmm, but true means that this insertelement is a part of buildvector. Do you want this?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9603

True means that we have a buildvector, but do not include this insertelement into the list of the vectorisable insertelements, i.e. treat it as an initial vector of the buildvector. Could you try this?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9607–9609

I mean, maybe return true here too, if we have at least 2 insertelements in the InsertElts vector?

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

9603

Ok, I see.

9607–9609

I've simplified this a bit more, removed boolean return type of function at all, since "at least 2 insertelements" condition has already been checking at the non-recursive wrapper findBuildAggregate().

[anton-afanasyev]

Interpret insertelements with undef/uncorrect/non-constant index as initial vector

[ABataev]

Would be good to test it using test suite at least, just in case we missed something.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

With an early check for insertelements in tryVectorizePair this can be turned to an assert rather than check and schedule cancellation.

[anton-afanasyev]

In D119623#3318056, @ABataev wrote:

Would be good to test it using test suite at least, just in case we missed something.

Sure, I've tested, it's ok (no changes).

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

Yes, done in D119679

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

Update this patch after D119679

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

Hmm, conversely, D119679 is based on this patch.

[ABataev]

LG

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

Ok, do not forget to remove this check then.

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

4026–4032

Sure, already removed.

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

static void inversePermutation(ArrayRef<unsigned> Indices,
   const unsigned E = Indices.size();
   Mask.resize(E, UndefMaskElem);
   for (unsigned I = 0; I < E; ++I)
     Mask[Indices[I]] = I;
 }
 
 /// \returns inserting index of InsertElement or InsertValue instruction,
 /// using Offset as base offset for index.
-static Optional<int> getInsertIndex(Value *InsertInst, unsigned Offset = 0) {
+static Optional<unsigned> getInsertIndex(Value *InsertInst,
+                                         unsigned Offset = 0) {
   int Index = Offset;
   if (auto *IE = dyn_cast<InsertElementInst>(InsertInst)) {
     if (auto *CI = dyn_cast<ConstantInt>(IE->getOperand(2))) {
       auto *VT = cast<FixedVectorType>(IE->getType());
       if (CI->getValue().uge(VT->getNumElements()))
-        return UndefMaskElem;
+        return None;
       Index *= VT->getNumElements();
       Index += CI->getZExtValue();
       return Index;
     }
-    if (isa<UndefValue>(IE->getOperand(2)))
-      return UndefMaskElem;
     return None;
   }
 
   auto *IV = cast<InsertValueInst>(InsertInst);
   Type *CurrentType = IV->getType();
   for (unsigned I : IV->indices()) {
     if (auto *ST = dyn_cast<StructType>(CurrentType)) {
       Index *= ST->getNumElements();
case Instruction::ExtractElement: {
       return;
     }
     case Instruction::InsertElement: {
       assert(ReuseShuffleIndicies.empty() && "All inserts should be unique");
 
       // Check that we have a buildvector and not a shuffle of 2 or more
       // different vectors.
       ValueSet SourceVectors;
-      int MinIdx = std::numeric_limits<int>::max();
       for (Value *V : VL) {
         SourceVectors.insert(cast<Instruction>(V)->getOperand(0));
-        Optional<int> Idx = *getInsertIndex(V);
-        if (!Idx || *Idx == UndefMaskElem)
-          continue;
-        MinIdx = std::min(MinIdx, *Idx);
+        if (getInsertIndex(V) == None) {
+          LLVM_DEBUG(dbgs() << "SLP: Gather of insertelement vectors with "
+                               "non-constant or undef index.\n");
+          newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx);
+          BS.cancelScheduling(VL, VL0);
+          return;
+        }
       }
 
       if (count_if(VL, [&SourceVectors](Value *V) {
             return !SourceVectors.contains(V);
           }) >= 2) {
         // Found 2nd source vector - cancel.
         LLVM_DEBUG(dbgs() << "SLP: Gather of insertelement vectors with "
                              "different source vectors.\n");
         newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx);
         BS.cancelScheduling(VL, VL0);
         return;
       }
 
       auto OrdCompare = [](const std::pair<int, int> &P1,
                            const std::pair<int, int> &P2) {
         return P1.first > P2.first;
       };
       PriorityQueue<std::pair<int, int>, SmallVector<std::pair<int, int>>,
                     decltype(OrdCompare)>
           Indices(OrdCompare);
       for (int I = 0, E = VL.size(); I < E; ++I) {
-        Optional<int> Idx = *getInsertIndex(VL[I]);
-        if (!Idx || *Idx == UndefMaskElem)
-          continue;
-        Indices.emplace(*Idx, I);
+        unsigned Idx = *getInsertIndex(VL[I]);
+        Indices.emplace(Idx, I);
       }
       OrdersType CurrentOrder(VL.size(), VL.size());
       bool IsIdentity = true;
       for (int I = 0, E = VL.size(); I < E; ++I) {
         CurrentOrder[Indices.top().second] = I;
         IsIdentity &= Indices.top().second == I;
         Indices.pop();
       }
case Instruction::InsertElement: {
         Mask.assign(NumElts, UndefMaskElem);
         std::iota(Mask.begin(), std::next(Mask.begin(), NumScalars), 0);
       }
       unsigned Offset = *getInsertIndex(VL0);
       bool IsIdentity = true;
       SmallVector<int> PrevMask(NumElts, UndefMaskElem);
       Mask.swap(PrevMask);
       for (unsigned I = 0; I < NumScalars; ++I) {
-        Optional<int> InsertIdx = getInsertIndex(VL[PrevMask[I]]);
-        if (!InsertIdx || *InsertIdx == UndefMaskElem)
-          continue;
-        DemandedElts.setBit(*InsertIdx);
-        IsIdentity &= *InsertIdx - Offset == I;
-        Mask[*InsertIdx - Offset] = I;
+        unsigned InsertIdx = *getInsertIndex(VL[PrevMask[I]]);
+        DemandedElts.setBit(InsertIdx);
+        IsIdentity &= InsertIdx - Offset == I;
+        Mask[InsertIdx - Offset] = I;
       }
       assert(Offset < NumElts && "Failed to find vector index offset");
 
       InstructionCost Cost = 0;
       Cost -= TTI->getScalarizationOverhead(SrcVecTy, DemandedElts,
                                             /*Insert*/ true, /*Extract*/ false);
 
       if (IsIdentity && NumElts != NumScalars && Offset % NumScalars != 0) {
for (ExternalUser &EU : ExternalUses) {
     // for extractelements, no need to add it again.
     if (isa<ExtractElementInst>(EU.Scalar))
       continue;
 
     // If found user is an insertelement, do not calculate extract cost but try
     // to detect it as a final shuffled/identity match.
     if (auto *VU = dyn_cast_or_null<InsertElementInst>(EU.User)) {
       if (auto *FTy = dyn_cast<FixedVectorType>(VU->getType())) {
-        Optional<int> InsertIdx = getInsertIndex(VU);
-        if (!InsertIdx || *InsertIdx == UndefMaskElem)
-          continue;
+        unsigned InsertIdx = *getInsertIndex(VU);
         auto *It = find_if(FirstUsers, [VU](Value *V) {
           return areTwoInsertFromSameBuildVector(VU,
                                                  cast<InsertElementInst>(V));
         });
         int VecId = -1;
         if (It == FirstUsers.end()) {
           VF.push_back(FTy->getNumElements());
           ShuffleMask.emplace_back(VF.back(), UndefMaskElem);
if (auto *VU = dyn_cast_or_null<InsertElementInst>(EU.User)) {
             Base = cast<InsertElementInst>(Base)->getOperand(0);
           }
           FirstUsers.push_back(VU);
           DemandedElts.push_back(APInt::getZero(VF.back()));
           VecId = FirstUsers.size() - 1;
         } else {
           VecId = std::distance(FirstUsers.begin(), It);
         }
-        int Idx = *InsertIdx;
-        ShuffleMask[VecId][Idx] = EU.Lane;
-        DemandedElts[VecId].setBit(Idx);
+        ShuffleMask[VecId][InsertIdx] = EU.Lane;
+        DemandedElts[VecId].setBit(InsertIdx);
         continue;
       }
     }
 
     // If we plan to rewrite the tree in a smaller type, we will need to sign
     // extend the extracted value back to the original type. Here, we account
     // for the extract and the added cost of the sign extend if needed.
     auto *VecTy = FixedVectorType::get(EU.Scalar->getType(), BundleWidth);
case Instruction::InsertElement: {
         std::iota(Mask.begin(), std::next(Mask.begin(), NumScalars), 0);
       }
       // Create InsertVector shuffle if necessary
       bool IsIdentity = true;
       SmallVector<int> PrevMask(NumElts, UndefMaskElem);
       Mask.swap(PrevMask);
       for (unsigned I = 0; I < NumScalars; ++I) {
         Value *Scalar = E->Scalars[PrevMask[I]];
-        Optional<int> InsertIdx = getInsertIndex(Scalar);
-        if (!InsertIdx || *InsertIdx == UndefMaskElem)
-          continue;
-        IsIdentity &= *InsertIdx - Offset == I;
-        Mask[*InsertIdx - Offset] = I;
+        unsigned InsertIdx = *getInsertIndex(Scalar);
+        IsIdentity &= InsertIdx - Offset == I;
+        Mask[InsertIdx - Offset] = I;
       }
       if (!IsIdentity || NumElts != NumScalars) {
         V = Builder.CreateShuffleVector(V, Mask);
         if (auto *I = dyn_cast<Instruction>(V)) {
           GatherShuffleSeq.insert(I);
           CSEBlocks.insert(I->getParent());
         }
       }
do {
     } else if (CurrentType->isSingleValueType()) {
       return AggregateSize;
     } else {
       return None;
     }
   } while (true);
 }
 
-static bool findBuildAggregate_rec(Instruction *LastInsertInst,
+static void findBuildAggregate_rec(Instruction *LastInsertInst,
                                    TargetTransformInfo *TTI,
                                    SmallVectorImpl<Value *> &BuildVectorOpds,
                                    SmallVectorImpl<Value *> &InsertElts,
                                    unsigned OperandOffset) {
   do {
     Value *InsertedOperand = LastInsertInst->getOperand(1);
-    Optional<int> OperandIndex = getInsertIndex(LastInsertInst, OperandOffset);
+    Optional<unsigned> OperandIndex =
+        getInsertIndex(LastInsertInst, OperandOffset);
     if (!OperandIndex)
-      return false;
+      return;
     if (isa<InsertElementInst>(InsertedOperand) ||
         isa<InsertValueInst>(InsertedOperand)) {
-      if (!findBuildAggregate_rec(cast<Instruction>(InsertedOperand), TTI,
-                                  BuildVectorOpds, InsertElts, *OperandIndex))
-        return false;
+      findBuildAggregate_rec(cast<Instruction>(InsertedOperand), TTI,
+                             BuildVectorOpds, InsertElts, *OperandIndex);
+
     } else {
       BuildVectorOpds[*OperandIndex] = InsertedOperand;
       InsertElts[*OperandIndex] = LastInsertInst;
     }
     LastInsertInst = dyn_cast<Instruction>(LastInsertInst->getOperand(0));
   } while (LastInsertInst != nullptr &&
            (isa<InsertValueInst>(LastInsertInst) ||
             isa<InsertElementInst>(LastInsertInst)) &&
            LastInsertInst->hasOneUse());
-  return true;
 }
 
 /// Recognize construction of vectors like
 ///  %ra = insertelement <4 x float> poison, float %s0, i32 0
 ///  %rb = insertelement <4 x float> %ra, float %s1, i32 1
 ///  %rc = insertelement <4 x float> %rb, float %s2, i32 2
 ///  %rd = insertelement <4 x float> %rc, float %s3, i32 3
 ///  starting from the last insertelement or insertvalue instruction.
assert((BuildVectorOpds.empty() && InsertElts.empty()) &&
          "Expected empty result vectors!");
 
   Optional<unsigned> AggregateSize = getAggregateSize(LastInsertInst);
   if (!AggregateSize)
     return false;
   BuildVectorOpds.resize(*AggregateSize);
   InsertElts.resize(*AggregateSize);
 
-  if (findBuildAggregate_rec(LastInsertInst, TTI, BuildVectorOpds, InsertElts,
-                             0)) {
-    llvm::erase_value(BuildVectorOpds, nullptr);
-    llvm::erase_value(InsertElts, nullptr);
-    if (BuildVectorOpds.size() >= 2)
-      return true;
-  }
+  findBuildAggregate_rec(LastInsertInst, TTI, BuildVectorOpds, InsertElts, 0);
 
   return false;
 }
 
 /// Try and get a reduction value from a phi node.
 ///
 /// Given a phi node \p P in a block \p ParentBB, consider possible reductions
 /// if they come from either \p ParentBB or a containing loop latch.

llvm/test/Transforms/SLPVectorizer/X86/insert-crash-index.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -slp-vectorizer < %s | FileCheck %s
+
+; These all crashing before patch
+
+define  <2 x i8> @negative_index(<2 x i8> %aux_vec, i8* %in) {
+; CHECK-LABEL: @negative_index(
+; CHECK-NEXT:    [[IN0:%.*]] = load i8, i8* [[IN:%.*]], align 4
+; CHECK-NEXT:    [[G1:%.*]] = getelementptr inbounds i8, i8* [[IN]], i64 1
+; CHECK-NEXT:    [[IN1:%.*]] = load i8, i8* [[G1]], align 4
+; CHECK-NEXT:    [[V0:%.*]] = insertelement <2 x i8> [[AUX_VEC:%.*]], i8 [[IN0]], i8 -1
+; CHECK-NEXT:    [[V1:%.*]] = insertelement <2 x i8> [[V0]], i8 [[IN1]], i64 1
+; CHECK-NEXT:    [[OUT:%.*]] = add <2 x i8> [[V1]], [[V1]]
+; CHECK-NEXT:    ret <2 x i8> [[OUT]]
+;
+  %in0 = load i8, i8* %in, align 4
+  %g1 = getelementptr inbounds i8, i8* %in, i64 1
+  %in1 = load i8, i8* %g1, align 4
+
+  %v0 = insertelement <2 x i8> %aux_vec, i8 %in0, i8 -1
+  %v1 = insertelement <2 x i8> %v0, i8 %in1, i64 1
+
+  %out = add <2 x i8> %v1, %v1
+  ret <2 x i8> %out
+}
+
+define  <2 x i8> @exceed_index(<2 x i8> %aux_vec, i8* %in) {
+; CHECK-LABEL: @exceed_index(
+; CHECK-NEXT:    [[IN0:%.*]] = load i8, i8* [[IN:%.*]], align 4
+; CHECK-NEXT:    [[G1:%.*]] = getelementptr inbounds i8, i8* [[IN]], i64 1
+; CHECK-NEXT:    [[IN1:%.*]] = load i8, i8* [[G1]], align 4
+; CHECK-NEXT:    [[V0:%.*]] = insertelement <2 x i8> [[AUX_VEC:%.*]], i8 [[IN0]], i8 2
+; CHECK-NEXT:    [[V1:%.*]] = insertelement <2 x i8> [[V0]], i8 [[IN1]], i64 1
+; CHECK-NEXT:    [[OUT:%.*]] = add <2 x i8> [[V1]], [[V1]]
+; CHECK-NEXT:    ret <2 x i8> [[OUT]]
+;
+  %in0 = load i8, i8* %in, align 4
+  %g1 = getelementptr inbounds i8, i8* %in, i64 1
+  %in1 = load i8, i8* %g1, align 4
+
+  %v0 = insertelement <2 x i8> %aux_vec, i8 %in0, i8 2
+  %v1 = insertelement <2 x i8> %v0, i8 %in1, i64 1
+
+  %out = add <2 x i8> %v1, %v1
+  ret <2 x i8> %out
+}
+
+define  <2 x i8> @poison_index(<2 x i8> %aux_vec, i8* %in) {
+; CHECK-LABEL: @poison_index(
+; CHECK-NEXT:    [[IN0:%.*]] = load i8, i8* [[IN:%.*]], align 4
+; CHECK-NEXT:    [[G1:%.*]] = getelementptr inbounds i8, i8* [[IN]], i64 1
+; CHECK-NEXT:    [[IN1:%.*]] = load i8, i8* [[G1]], align 4
+; CHECK-NEXT:    [[V0:%.*]] = insertelement <2 x i8> [[AUX_VEC:%.*]], i8 [[IN0]], i8 poison
+; CHECK-NEXT:    [[V1:%.*]] = insertelement <2 x i8> [[V0]], i8 [[IN1]], i64 1
+; CHECK-NEXT:    [[OUT:%.*]] = add <2 x i8> [[V1]], [[V1]]
+; CHECK-NEXT:    ret <2 x i8> [[OUT]]
+;
+  %in0 = load i8, i8* %in, align 4
+  %g1 = getelementptr inbounds i8, i8* %in, i64 1
+  %in1 = load i8, i8* %g1, align 4
+
+  %v0 = insertelement <2 x i8> %aux_vec, i8 %in0, i8 poison
+  %v1 = insertelement <2 x i8> %v0, i8 %in1, i64 1
+
+  %out = add <2 x i8> %v1, %v1
+  ret <2 x i8> %out
+}