[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

 /// %9 = extractelement <4 x i8> %5, i32 3
 /// %ins4 = insertelement <4 x i8> %ins3, i8 %9, i32 3
 /// ret <4 x i8> %ins4
 /// InstCombiner transforms this into a shuffle and vector mul
 /// Mask will return the Shuffle Mask equivalent to the extracted elements.
 /// TODO: Can we split off and reuse the shuffle mask detection from
 /// TargetTransformInfo::getInstructionThroughput?
 static Optional<TargetTransformInfo::ShuffleKind>
-isShuffle(ArrayRef<Value *> VL, SmallVectorImpl<int> &Mask) {
+isFixedVectorShuffle(ArrayRef<Value *> VL, SmallVectorImpl<int> &Mask) {
   auto *EI0 = cast<ExtractElementInst>(VL[0]);
+  if (isa<ScalableVectorType>(EI0->getVectorOperandType()))
+    return None;
   unsigned Size =
       cast<FixedVectorType>(EI0->getVectorOperandType())->getNumElements();
   Value *Vec1 = nullptr;
   Value *Vec2 = nullptr;
   enum ShuffleMode { Unknown, Select, Permute };
   ShuffleMode CommonShuffleMode = Unknown;
   for (unsigned I = 0, E = VL.size(); I < E; ++I) {
     auto *EI = cast<ExtractElementInst>(VL[I]);
if (E->State == TreeEntry::NeedToGather) {
     if (E->getOpcode() == Instruction::ExtractElement && allSameType(VL) &&
         allSameBlock(VL) &&
         !isa<ScalableVectorType>(
             cast<ExtractElementInst>(E->getMainOp())->getVectorOperandType())) {
       // Check that gather of extractelements can be represented as just a
       // shuffle of a single/two vectors the scalars are extracted from.
       SmallVector<int> Mask;
       Optional<TargetTransformInfo::ShuffleKind> ShuffleKind =
-          isShuffle(VL, Mask);
+          isFixedVectorShuffle(VL, Mask);
       if (ShuffleKind.hasValue()) {
         // Found the bunch of extractelement instructions that must be gathered
         // into a vector and can be represented as a permutation elements in a
         // single input vector or of 2 input vectors.
         InstructionCost Cost =
             computeExtractCost(VL, VecTy, *ShuffleKind, Mask, *TTI);
         AdjustExtractsCost(Cost, /*IsGather=*/true);
         if (NeedToShuffleReuses)
bool BoUpSLP::isFullyVectorizableTinyTree() const {
   if (VectorizableTree[0]->State == TreeEntry::Vectorize &&
       (allConstant(VectorizableTree[1]->Scalars) ||
        isSplat(VectorizableTree[1]->Scalars) ||
        (VectorizableTree[1]->State == TreeEntry::NeedToGather &&
         VectorizableTree[1]->Scalars.size() <
             VectorizableTree[0]->Scalars.size()) ||
        (VectorizableTree[1]->State == TreeEntry::NeedToGather &&
         VectorizableTree[1]->getOpcode() == Instruction::ExtractElement &&
-        isShuffle(VectorizableTree[1]->Scalars, Mask))))
+        isFixedVectorShuffle(VectorizableTree[1]->Scalars, Mask))))
     return true;
 
   // Gathering cost would be too much for tiny trees.
   if (VectorizableTree[0]->State == TreeEntry::NeedToGather ||
       VectorizableTree[1]->State == TreeEntry::NeedToGather)
     return false;
 
   return true;

 bool SLPVectorizerPass::vectorizeInsertElementInst(InsertElementInst *IEI,
                                                    BasicBlock *BB, BoUpSLP &R) {
   SmallVector<Value *, 16> BuildVectorInsts;
   SmallVector<Value *, 16> BuildVectorOpds;
   SmallVector<int> Mask;
   if (!findBuildAggregate(IEI, TTI, BuildVectorOpds, BuildVectorInsts) ||
       (llvm::all_of(BuildVectorOpds,
                     [](Value *V) { return isa<ExtractElementInst>(V); }) &&
-       isShuffle(BuildVectorOpds, Mask)))
+       isFixedVectorShuffle(BuildVectorOpds, Mask)))
     return false;
 
   LLVM_DEBUG(dbgs() << "SLP: array mappable to vector: " << *IEI << "\n");
   return tryToVectorizeList(BuildVectorInsts, R);
 }
 
 bool SLPVectorizerPass::vectorizeSimpleInstructions(
     SmallVectorImpl<Instruction *> &Instructions, BasicBlock *BB, BoUpSLP &R,

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>)