[SLP] Fix assert from non-constant index in insertelement

[bcahoon]

Summary

A call to getInsertIndex() in getTreeCost() is returning None,
which causes an assert because a non-constant index value for
insertelement was not expected. This case occurs when the
insertelement index value is defined with a PHI.

I'm not sure if this is the best way to fix this issue, but it does
fix the assert.

[anton-afanasyev]

This is a partial revert of rG954ea0f044e0. It looks like there is still an issue with cost computation here, but I'm to fix it in a separate patch and to backport both patches to 14.0.0. Therefore looks good to me. Add @ABataev.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

I think we need an extra check here. If the index is undefined - continue, but if the index is not constant, need to calculate the cost as extractelement+insertelement.

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

This is just a reversion to previous state, we can proceed with this and commit a separate patch for correct check. There is also another check we need here: that EU.Scalar is used as a first operand of insertelement.

5921–5922

As for undef index check, I'm not sure it should be processed specially: we still generate extractelement for EU.Scalar at vectorizeTree() stage.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

No, we don't need this check here.
I was not quite correct in my previous comment. No need to to calculate the cost as extractelement+insertelement, just extractelement.
I mean, we need to have something like this:

Optional<unsigned> Idx = getInsertIndex(VU);
if (Idx) {
  // Process as shuffle.
  ...
} else if (isa<Constant>(VY->getOperand(2))
  continue;

5921–5922

Yes, but it will optimized out by InstCombiner, since the user with undef index is just an undef.

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

No, we don't need this check here.

What if EU.Scalar is used a second operand?

5921–5922

Sure, it will optimized in that case, but why do we need this special check for insertelement here? Why not for mul i32 ..., poison, for instance? If we want to check that the value extracted is used by poison only, we can do this generally in special place.

5921–5922

I propose to proceed with current fix and then go on further.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Still need to extract it and count the cost of extract.

5921–5922

Generally speaking, it would be good to have such checks, we just can't check everything because of the compile time and amount of work and rely on InstCombiner here, that it transforms such instructions to undefs before SLP pass. We just have some artificial tests, which has undef indices, this check is mostly for such test cases.

5921–5922

We already underestimate the cost in many cases, would be good to have the proper fix.

[anton-afanasyev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Sure. So this is another, main branch. So we need something like EU.Scalar == VU.getOperand(0) check here.

5921–5922

Ok, I see. So do you believe @bcahoon could commit this and then I add other checks?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

It means that the scalar is an insertelement instruction, i.e. has vector type, right? It is checked already in line 5908.

5921–5922

For the 14.0 branch? I would suggest following more conservative solution, probably:

Optional<unsigned> Idx = getInsertIndex(VU);
if (Idx) {
  // Process as shuffle.
  ...
}

i.e. do not ignore it but instead count as extractelement cost

[bcahoon]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Thanks for the reviews so far. I'm trying to figure out the next steps to fix the assert due to dereferencing the result of getInsertIndex(VU) when it returns None. What's involved with the comment,

// Process as shuffle.

Is that the existing code that appears in lines 5923-5949?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Yes

[bcahoon]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Isn't that the behavior of the patch? The existing code is skipped if (!Idx), otherwise Idx contains a valid value, so execute the existing code. Does something else need to be added?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

No, the patch does continue, i.e. skips everything, but instead it should fallback to the cost for extractelement.

[bcahoon]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Ah, I think I understand. When getInsertValue returns None, the lines 5962-5970 should still occurs rather than being skipped. Those lines are performing the "cost for extractelement"? I will try that.

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

5921–5922

Yep, you got it right.

[bcahoon]

Changes based upon review. Don't just skip over insertelements with a
non-constant index. Instead, fall through to include additional cost
calculation.

[ABataev]

LG

[anton-afanasyev]

I'm to backport this: https://github.com/llvm/llvm-project/issues/54013

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

for (ExternalUser &EU : ExternalUses) {
     // Already counted the cost for external uses when tried to adjust the cost
     // 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())) {
-        unsigned InsertIdx = *getInsertIndex(VU);
+        Optional<unsigned> InsertIdx = getInsertIndex(VU);
+        if (InsertIdx) {
           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);
             // Find the insertvector, vectorized in tree, if any.
             Value *Base = VU;
             while (isa<InsertElementInst>(Base)) {
               // Build the mask for the vectorized insertelement instructions.
               if (const TreeEntry *E = getTreeEntry(Base)) {
                 VU = cast<InsertElementInst>(Base);
                 do {
                   int Idx = E->findLaneForValue(Base);
                   ShuffleMask.back()[Idx] = Idx;
                   Base = cast<InsertElementInst>(Base)->getOperand(0);
                 } while (E == getTreeEntry(Base));
                 break;
               }
               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);
           }
-        ShuffleMask[VecId][InsertIdx] = EU.Lane;
-        DemandedElts[VecId].setBit(InsertIdx);
+          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);
     auto *ScalarRoot = VectorizableTree[0]->Scalars[0];
     if (MinBWs.count(ScalarRoot)) {
       auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot].first);

llvm/test/Transforms/SLPVectorizer/slp-variable-insertelement.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -slp-vectorizer -slp-max-vf=2 -slp-min-reg-size=32 -S < %s | FileCheck %s
+
+; It is possible to compute the tree cost for an insertelement that does not
+; have a constant index when the index is a PHI. Check if getInsertIndex
+; returns None.
+
+define void @test() local_unnamed_addr {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    unreachable
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ poison, [[FOR_BODY]] ]
+; CHECK-NEXT:    [[J:%.*]] = phi i32 [ poison, [[ENTRY]] ], [ 0, [[FOR_BODY]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i32> poison, i32 poison, i32 [[I]]
+; CHECK-NEXT:    br i1 poison, label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup:
+  unreachable
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ poison, %for.body ]
+  %j = phi i32 [ poison, %entry ], [ 0, %for.body ]
+  %0 = insertelement <4 x i32> poison, i32 poison, i32 %i
+  br i1 poison, label %for.cond.cleanup, label %for.body
+}