[llvm] [SLP] Avoid crash in computeExtractCost (PR #93188)

[Björn Pettersson via llvm-commits]

https://github.com/bjope created https://github.com/llvm/llvm-project/pull/93188

For a downstream target we ended up in a situation with assertion failures in ShuffleCostEstimator::computeExtractCost.

Input IR looked like this:

define void @foo(ptr %p0, <64 x i32> %vec) {
  %p1 = getelementptr i32, ptr %p0, i16 1
  %p2 = getelementptr i32, ptr %p0, i16 2
  %p3 = getelementptr i32, ptr %p0, i16 3
  %p4 = getelementptr i32, ptr %p0, i16 4
  %p5 = getelementptr i32, ptr %p0, i16 5
  %p6 = getelementptr i32, ptr %p0, i16 6
  %p7 = getelementptr i32, ptr %p0, i16 7
  %elt = extractelement <64 x i32> %vec, i32 0
  store i32 %elt, ptr %p0
  store i32 %elt, ptr %p1
  store i32 %elt, ptr %p2
  store i32 %elt, ptr %p3
  store i32 %elt, ptr %p4
  store i32 %elt, ptr %p5
  store i32 %elt, ptr %p6
  store i32 %elt, ptr %p7
  ret void
}

And the scenario was like this:
- VL and Mask has 8 elements at entry to computeExtractCost.
- NumParts is 2 (v8i32 is not legal, but v4i32 is legal).
- NumElts is calculated as 64 (given by extractelement <64 x i32>).
- NumSrcRegs is calculated and is set to 1 (v64i32 is legal).
- EltsPerVector is calculated as 64 (given by NumElts/NumSrcRegs).
- Assertion failure happens when doing ArrayRef<int> MaskSlice = Mask.slice(Part * EltsPerVector, (Part == NumParts - 1 && Mask.size() % EltsPerVector != 0) ? Mask.size() % EltsPerVector : EltsPerVector); since EltsPerVector is larger than Mask.size() already for Part==0.

This patch resolved the issue by making sure that we slice up Mask in at most EltsPerVector pieces until we have covered the full Mask. When we have covered all elements in Mask we break the loop.

Haven't been able to reproduce this scenario for any in-tree target. So unfortunately there is no regression test included in the patch.

>From dac96ed216f352a06da6db4c4fd7c5aee41765d9 Mon Sep 17 00:00:00 2001
From: Bjorn Pettersson <bjorn.a.pettersson at ericsson.com>
Date: Thu, 23 May 2024 15:04:11 +0200
Subject: [PATCH] [SLP] Avoid crash in computeExtractCost

For a downstream target we ended up in a situation with assertion
failures in ShuffleCostEstimator::computeExtractCost.

Input IR looked like this:

define void @foo(ptr %p0, <64 x i32> %vec) {
  %p1 = getelementptr i32, ptr %p0, i16 1
  %p2 = getelementptr i32, ptr %p0, i16 2
  %p3 = getelementptr i32, ptr %p0, i16 3
  %p4 = getelementptr i32, ptr %p0, i16 4
  %p5 = getelementptr i32, ptr %p0, i16 5
  %p6 = getelementptr i32, ptr %p0, i16 6
  %p7 = getelementptr i32, ptr %p0, i16 7
  %elt = extractelement <64 x i32> %vec, i32 0
  store i32 %elt, ptr %p0
  store i32 %elt, ptr %p1
  store i32 %elt, ptr %p2
  store i32 %elt, ptr %p3
  store i32 %elt, ptr %p4
  store i32 %elt, ptr %p5
  store i32 %elt, ptr %p6
  store i32 %elt, ptr %p7
  ret void
}

And the scenario was like this:
- VL and Mask has 8 elements at entry to computeExtractCost.
- NumParts is 2 (v8i32 is not legal, but v4i32 is legal).
- NumElts is calculated as 64 (given by extractelement <64 x i32>).
- NumSrcRegs is calculated and is set to 1 (v64i32 is legal).
- EltsPerVector is calculated as 64 (given by NumElts/NumSrcRegs).
- Assertion failure happens when doing
      ArrayRef<int> MaskSlice =
          Mask.slice(Part * EltsPerVector,
                     (Part == NumParts - 1 && Mask.size() % EltsPerVector != 0)
                         ? Mask.size() % EltsPerVector
                         : EltsPerVector);
  since EltsPerVector is larger than Mask.size() already for Part==0.

This patch resolved the issue by making sure that we slice up Mask
in at most EltsPerVector pieces until we have covered the full Mask.
When we have covered all elements in Mask we break the loop.

Haven't been able to reproduce this scenario for any in-tree target.
So unfortunately there is no regression test included in the patch.
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 08ecbe304429e..f25ddf6bd8082 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -8318,11 +8318,13 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
     for (unsigned Part = 0; Part < NumParts; ++Part) {
       if (!ShuffleKinds[Part])
         continue;
-      ArrayRef<int> MaskSlice =
-          Mask.slice(Part * EltsPerVector,
-                     (Part == NumParts - 1 && Mask.size() % EltsPerVector != 0)
-                         ? Mask.size() % EltsPerVector
-                         : EltsPerVector);
+      unsigned SliceStart = Part * EltsPerVector;
+      if (SliceStart >= Mask.size())
+        break;
+      unsigned SliceSize = (SliceStart + EltsPerVector) > Mask.size()
+                               ? Mask.size() - SliceStart
+                               : EltsPerVector;
+      ArrayRef<int> MaskSlice = Mask.slice(SliceStart, SliceSize);
       SmallVector<int> SubMask(EltsPerVector, PoisonMaskElem);
       copy(MaskSlice, SubMask.begin());
       std::optional<TTI::ShuffleKind> RegShuffleKind =