[PATCH] D110640: [SLPVectorizer] Fix crash in isShuffle with scalable vectors

[Kerry McLaughlin via Phabricator via llvm-commits]

kmclaughlin created this revision.
kmclaughlin added reviewers: CarolineConcatto, ABataev, sdesmalen, david-arm.
Herald added a subscriber: hiraditya.
kmclaughlin requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

D104809 <https://reviews.llvm.org/D104809> changed `buildTree_rec` to check for extract element instructions
with scalable types. However, if the extract is extended or truncated,
these changes do not apply and we assert later on in isShuffle(), which
attempts to cast the type of the extract to FixedVectorType.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D110640

Files:
  llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
  llvm/test/Transforms/SLPVectorizer/AArch64/scalable-vector.ll


Index: llvm/test/Transforms/SLPVectorizer/AArch64/scalable-vector.ll
===================================================================
--- llvm/test/Transforms/SLPVectorizer/AArch64/scalable-vector.ll
+++ llvm/test/Transforms/SLPVectorizer/AArch64/scalable-vector.ll
@@ -138,5 +138,62 @@
   ret  <vscale x 4 x i8> %ins4
 }
 
+define void @sext_scalable_extractelement() {
+; CHECK-LABEL: @sext_scalable_extractelement(
+; CHECK-NEXT:    [[X0:%.*]] = extractelement <vscale x 2 x i32> undef, i32 undef
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[X0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i64, i64* undef, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <vscale x 2 x i32> undef, i32 undef
+; CHECK-NEXT:    [[TMP4:%.*]] = sext i32 [[TMP3]] to i64
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds i64, i64* undef, i64 [[TMP4]]
+; CHECK-NEXT:    ret void
+;
+  %x0 = extractelement <vscale x 2 x i32> undef, i32 undef
+  %1 = sext i32 %x0 to i64
+  %2 = getelementptr inbounds i64, i64* undef, i64 %1
+  %3 = extractelement <vscale x 2 x i32> undef, i32 undef
+  %4 = sext i32 %3 to i64
+  %5 = getelementptr inbounds i64, i64* undef, i64 %4
+  ret void
+}
+
+define void @zext_scalable_extractelement() {
+; CHECK-LABEL: @zext_scalable_extractelement(
+; CHECK-NEXT:    [[X0:%.*]] = extractelement <vscale x 2 x i32> undef, i32 undef
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[X0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i64, i64* undef, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <vscale x 2 x i32> undef, i32 undef
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[TMP3]] to i64
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds i64, i64* undef, i64 [[TMP4]]
+; CHECK-NEXT:    ret void
+;
+  %x0 = extractelement <vscale x 2 x i32> undef, i32 undef
+  %1 = zext i32 %x0 to i64
+  %2 = getelementptr inbounds i64, i64* undef, i64 %1
+  %3 = extractelement <vscale x 2 x i32> undef, i32 undef
+  %4 = zext i32 %3 to i64
+  %5 = getelementptr inbounds i64, i64* undef, i64 %4
+  ret void
+}
+
+define void @trunc_scalable_extractelement() {
+; CHECK-LABEL: @trunc_scalable_extractelement(
+; CHECK-NEXT:    [[X0:%.*]] = extractelement <vscale x 2 x i64> undef, i32 undef
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[X0]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, i32* undef, i32 [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <vscale x 2 x i64> undef, i32 undef
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i64 [[TMP3]] to i32
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds i32, i32* undef, i32 [[TMP4]]
+; CHECK-NEXT:    ret void
+;
+  %x0 = extractelement <vscale x 2 x i64> undef, i32 undef
+  %1 = trunc i64 %x0 to i32
+  %2 = getelementptr inbounds i32, i32* undef, i32 %1
+  %3 = extractelement <vscale x 2 x i64> undef, i32 undef
+  %4 = trunc i64 %3 to i32
+  %5 = getelementptr inbounds i32, i32* undef, i32 %4
+  ret void
+}
+
 declare <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>*, i32 immarg, <vscale x 16 x i1>, <vscale x 16 x i8>)
 declare void @llvm.masked.store.nxv16i8.p0nxv16i8(<vscale x 16 x i8>, <vscale x 16 x i8>*, i32 immarg, <vscale x 16 x i1>)
Index: llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
===================================================================
--- llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -3191,6 +3191,20 @@
     return;
   }
 
+  if (S.getOpcode() == Instruction::SExt ||
+      S.getOpcode() == Instruction::ZExt ||
+      S.getOpcode() == Instruction::Trunc) {
+    auto *ExtInst = cast<Instruction>(S.OpValue);
+    Instruction *Op = dyn_cast<Instruction>(ExtInst->getOperand(0));
+    if (Op && Op->getOpcode() == Instruction::ExtractElement &&
+        isa<ScalableVectorType>(
+            cast<ExtractElementInst>(Op)->getVectorOperandType())) {
+      LLVM_DEBUG(dbgs() << "SLP: Gathering due to scalable vector type.\n");
+      newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx);
+      return;
+    }
+  }
+
   // Don't handle vectors.
   if (S.OpValue->getType()->isVectorTy() &&
       !isa<InsertElementInst>(S.OpValue)) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D110640.375633.patch
Type: text/x-patch
Size: 4212 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210928/9e960a43/attachment.bin>