[SLPVectorizer] Fix crash in isShuffle with scalable vectors

[kmclaughlin]

Summary

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.

[kmclaughlin]

Hi @ABataev, I've added a fix for this assertion failure in the same place as you suggested on a previous patch, D104809. Since this is quite similar and there would now be two fixes in buildTree_rec for the same issue, I was wondering if you think it's worth revisiting the fix @CarolineConcatto first proposed on that patch? Her first revision added the phiNodeHasScalableOp & basicBlockHasScalableOp functions to avoid vectorizing early if there are any instructions found with scalable types.

[ABataev]

In D110640#3028002, @kmclaughlin wrote:

Hi @ABataev, I've added a fix for this assertion failure in the same place as you suggested on a previous patch, D104809. Since this is quite similar and there would now be two fixes in buildTree_rec for the same issue, I was wondering if you think it's worth revisiting the fix @CarolineConcatto first proposed on that patch? Her first revision added the phiNodeHasScalableOp & basicBlockHasScalableOp functions to avoid vectorizing early if there are any instructions found with scalable types.

I think the current implementation is still good. Also, isShuffle is called in several places, which one causes the crash?

[kmclaughlin]

In D110640#3028021, @ABataev wrote:
I think the current implementation is still good. Also, isShuffle is called in several places, which one causes the crash?

It's the isShuffle call in isFullyVectorizableTinyTree that caused the crash here, which was reached through vectorizeGEPIndices->tryToVectorizeList.

[ABataev]

In D110640#3029994, @kmclaughlin wrote:

In D110640#3028021, @ABataev wrote:
I think the current implementation is still good. Also, isShuffle is called in several places, which one causes the crash?

It's the isShuffle call in isFullyVectorizableTinyTree that caused the crash here, which was reached through vectorizeGEPIndices->tryToVectorizeList.

Just check for scalable vectors before this isShuffle function call.

[sdesmalen]

In D110640#3029995, @ABataev wrote:

Just check for scalable vectors before this isShuffle function call.

Alternatively, do we want to rename isShuffle to isFixedVectorShuffle and return None if the extract-element's vector operand type is scalable?

[ABataev]

In D110640#3030337, @sdesmalen wrote:

In D110640#3029995, @ABataev wrote:

Just check for scalable vectors before this isShuffle function call.

Alternatively, do we want to rename isShuffle to isFixedVectorShuffle and return None if the extract-element's vector operand type is scalable?

Yep, that's a good option too.

[kmclaughlin]

• Removed changes to buildTree_rec.
• Renamed isShuffle to isFixedVectorShuffle and changed this to return None if the type of the extract element instruction is scalable.

[ABataev]

LG

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

void BoUpSLP::buildTree_rec(ArrayRef<Value *> VL, unsigned Depth,
   if (S.getOpcode() == Instruction::ExtractElement &&
       isa<ScalableVectorType>(
           cast<ExtractElementInst>(S.OpValue)->getVectorOperandType())) {
     LLVM_DEBUG(dbgs() << "SLP: Gathering due to scalable vector type.\n");
     newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx);
     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)) {
     LLVM_DEBUG(dbgs() << "SLP: Gathering due to vector type.\n");
     newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx);
     return;
   }
 

llvm/test/Transforms/SLPVectorizer/AArch64/scalable-vector.ll

;
   %y2y2 = mul i8 %y2, %y2
   %ins1 = insertelement <vscale x 4 x i8> poison, i8 %x0x0, i32 0
   %ins2 = insertelement <vscale x 4 x i8> %ins1, i8 %x3x3, i32 1
   %ins3 = insertelement <vscale x 4 x i8> %ins2, i8 %y1y1, i32 2
   %ins4 = insertelement <vscale x 4 x i8> %ins3, i8 %y2y2, i32 3
   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>)