From 26fec4e845a833214d2b0036ea90f1b1c2418d50 Mon Sep 17 00:00:00 2001
From: Alexey Bataev <a.bataev@outlook.com>
Date: Tue, 3 Jan 2023 08:44:51 -0800
Subject: [PATCH] [SLP]Fix crash on casting non-instruction extractelement.

Need to check if the extractelement operation is an extraction before
trying to move it around the buildblocks to avoid crash on cast.
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp     |  7 ++++---
 .../SLPVectorizer/AArch64/external-non-inst-use.ll  | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/SLPVectorizer/AArch64/external-non-inst-use.ll

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index f852a21..f11dc59 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -9556,7 +9556,7 @@ Value *BoUpSLP::vectorizeTree(ExtraValueToDebugLocsMap &ExternallyUsedValues,
   DenseMap<Value *, InsertElementInst *> VectorToInsertElement;
   // Maps extract Scalar to the corresponding extractelement instruction in the
   // basic block. Only one extractelement per block should be emitted.
-  DenseMap<Value *, DenseMap<BasicBlock *, Value *>> ScalarToEEs;
+  DenseMap<Value *, DenseMap<BasicBlock *, Instruction *>> ScalarToEEs;
   // Extract all of the elements with the external uses.
   for (const auto &ExternalUse : ExternalUses) {
     Value *Scalar = ExternalUse.Scalar;
@@ -9588,7 +9588,7 @@ Value *BoUpSLP::vectorizeTree(ExtraValueToDebugLocsMap &ExternallyUsedValues,
           // current block.
           auto EEIt = It->second.find(Builder.GetInsertBlock());
           if (EEIt != It->second.end()) {
-            auto *I = cast<Instruction>(EEIt->second);
+            Instruction *I = EEIt->second;
             if (Builder.GetInsertPoint() != Builder.GetInsertBlock()->end() &&
                 Builder.GetInsertPoint()->comesBefore(I))
               I->moveBefore(&*Builder.GetInsertPoint());
@@ -9603,7 +9603,8 @@ Value *BoUpSLP::vectorizeTree(ExtraValueToDebugLocsMap &ExternallyUsedValues,
           } else {
             Ex = Builder.CreateExtractElement(Vec, Lane);
           }
-          ScalarToEEs[Scalar].try_emplace(Builder.GetInsertBlock(), Ex);
+          if (auto *I = dyn_cast<Instruction>(Ex))
+            ScalarToEEs[Scalar].try_emplace(Builder.GetInsertBlock(), I);
         }
         // The then branch of the previous if may produce constants, since 0
         // operand might be a constant.
diff --git a/llvm/test/Transforms/SLPVectorizer/AArch64/external-non-inst-use.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/external-non-inst-use.ll
new file mode 100644
index 0000000..d4e3fb3
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/AArch64/external-non-inst-use.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes="slp-vectorizer" -mtriple=aarch64--linux-gnu -S %s | FileCheck %s
+
+define i16 @foo(ptr %p1, ptr %p2) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 0, ptr [[P1:%.*]], align 1
+; CHECK-NEXT:    [[CONST_MAT:%.*]] = or i32 0, 0
+; CHECK-NEXT:    store <2 x i32> zeroinitializer, ptr [[P2:%.*]], align 1
+; CHECK-NEXT:    ret i16 0
+;
+entry:
+  %const1 = bitcast i32 0 to i32
+  %const = bitcast i32 0 to i32
+  store i32 %const, ptr %p1, align 1
+  %const_mat = or i32 %const, 0
+  %v_332.imagp = getelementptr { i32, i32 }, ptr %p2, i32 0, i32 1
+  store i32 %const, ptr %p2, align 1
+  store i32 %const1, ptr %v_332.imagp, align 1
+  ret i16 0
+}
-- 
2.7.4