[SLP]Fix a crash for gather node with instructions from different bbs,

if cost threshold is very low.

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -9275,11 +9275,16 @@ bool BoUpSLP::isTreeTinyAndNotFullyVectorizable(bool ForReduction) const {
 
   // Check if any of the gather node forms an insertelement buildvector
   // somewhere.
-  if (any_of(VectorizableTree, [](const std::unique_ptr<TreeEntry> &TE) {
+  bool IsAllowedSingleBVNode =
+      VectorizableTree.size() > 1 ||
+      (VectorizableTree.size() == 1 && VectorizableTree.front()->getOpcode() &&
+       allSameBlock(VectorizableTree.front()->Scalars));
+  if (any_of(VectorizableTree, [&](const std::unique_ptr<TreeEntry> &TE) {
         return TE->State == TreeEntry::NeedToGather &&
-               all_of(TE->Scalars, [](Value *V) {
+               all_of(TE->Scalars, [&](Value *V) {
                  return isa<ExtractElementInst, UndefValue>(V) ||
-                        (!V->hasNUsesOrMore(UsesLimit) &&
+                        (IsAllowedSingleBVNode &&
+                         !V->hasNUsesOrMore(UsesLimit) &&
                          any_of(V->users(), [](User *U) {
                            return isa<InsertElementInst>(U);
                          }));

llvm/test/Transforms/SLPVectorizer/X86/gather-nodes-different-bb.ll

@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=slp-vectorizer -mtriple=x86_64-unknown-linux -mattr="-avx512pf,+avx512f,+avx512bw" -slp-threshold=-100 < %s | FileCheck %s
+
+define i1 @foo(i32 %a) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = sub nsw i32 0, [[A:%.*]]
+; CHECK-NEXT:    br label [[BB4:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[LOCAL:%.*]] = sub nsw i32 0, 0
+; CHECK-NEXT:    [[INS1:%.*]] = insertelement <2 x i32> poison, i32 [[TMP0]], i32 0
+; CHECK-NEXT:    [[ADD:%.*]] = icmp eq i32 [[TMP0]], [[LOCAL]]
+; CHECK-NEXT:    ret i1 [[ADD]]
+;
+entry:
+  %0 = sub nsw i32 0, %a
+  br label %bb1
+
+bb1:
+  %local = sub nsw i32 0, 0
+  %ins1 = insertelement <2 x i32> poison, i32 %0, i32 0
+  %add = icmp eq i32 %0, %local
+  ret i1 %add
+}