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

[Brendon Cahoon via llvm-commits]

Author: Brendon Cahoon
Date: 2022-02-22T15:57:14-06:00
New Revision: 3cc15e2cb657f3a814c8bd482bb7108782561abd

URL: https://github.com/llvm/llvm-project/commit/3cc15e2cb657f3a814c8bd482bb7108782561abd
DIFF: https://github.com/llvm/llvm-project/commit/3cc15e2cb657f3a814c8bd482bb7108782561abd.diff

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

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.

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

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

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index bdc6a33eb2278..4ea8a77583cc9 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -5919,39 +5919,41 @@ InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
     // 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);
-        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;
+        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);
             }
-            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);
           }
-          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);
+          continue;
         }
-        ShuffleMask[VecId][InsertIdx] = EU.Lane;
-        DemandedElts[VecId].setBit(InsertIdx);
-        continue;
       }
     }
 

diff  --git a/llvm/test/Transforms/SLPVectorizer/slp-variable-insertelement.ll b/llvm/test/Transforms/SLPVectorizer/slp-variable-insertelement.ll
new file mode 100644
index 0000000000000..d97bfe59e55c6
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/slp-variable-insertelement.ll
@@ -0,0 +1,31 @@
+; 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
+}