[llvm] 8d7a6e2 - [SLP]Fix a crash for gather node with instructions from different bbs,

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2024-03-21T08:03:06-07:00
New Revision: 8d7a6e2fd8beadc54a2e54fa361d35c66fdb40a4

URL: https://github.com/llvm/llvm-project/commit/8d7a6e2fd8beadc54a2e54fa361d35c66fdb40a4
DIFF: https://github.com/llvm/llvm-project/commit/8d7a6e2fd8beadc54a2e54fa361d35c66fdb40a4.diff

LOG: [SLP]Fix a crash for gather node with instructions from different bbs,
if cost threshold is very low.

Added: 
    llvm/test/Transforms/SLPVectorizer/X86/gather-nodes-different-bb.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 4853c2006fea06..36b446962c4a63 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/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);
                          }));

diff  --git a/llvm/test/Transforms/SLPVectorizer/X86/gather-nodes-
diff erent-bb.ll b/llvm/test/Transforms/SLPVectorizer/X86/gather-nodes-
diff erent-bb.ll
new file mode 100644
index 00000000000000..2acbe89b97757f
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/gather-nodes-
diff erent-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
+}