[llvm] e664dea - [SLP]Fix write-after-bounds.

[Alexey Bataev via llvm-commits]

Author: Alexey Bataev
Date: 2022-09-21T08:00:15-07:00
New Revision: e664dea1821ab1277e62f0b4074fb02867636e6e

URL: https://github.com/llvm/llvm-project/commit/e664dea1821ab1277e62f0b4074fb02867636e6e
DIFF: https://github.com/llvm/llvm-project/commit/e664dea1821ab1277e62f0b4074fb02867636e6e.diff

LOG: [SLP]Fix write-after-bounds.

Mask might be larger than the NumElts-OffsetBeg, need to use actual
indices to avoid acces out of bounds.

Added: 
    llvm/test/Transforms/SLPVectorizer/X86/buildvector-insert-mask-size.ll

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 9eb3fe3f3e335..9aad015bb79e2 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -6340,14 +6340,17 @@ InstructionCost BoUpSLP::getEntryCost(const TreeEntry *E,
 
       unsigned NumOfParts = TTI->getNumberOfParts(SrcVecTy);
 
+      SmallVector<int> InsertMask(NumElts, UndefMaskElem);
       unsigned OffsetBeg = *getInsertIndex(VL.front());
       unsigned OffsetEnd = OffsetBeg;
-      for (Value *V : VL.drop_front()) {
+      InsertMask[OffsetBeg] = 0;
+      for (auto [I, V] : enumerate(VL.drop_front())) {
         unsigned Idx = *getInsertIndex(V);
         if (OffsetBeg > Idx)
           OffsetBeg = Idx;
         else if (OffsetEnd < Idx)
           OffsetEnd = Idx;
+        InsertMask[Idx] = I + 1;
       }
       unsigned VecScalarsSz = PowerOf2Ceil(NumElts);
       if (NumOfParts > 0)
@@ -6412,8 +6415,6 @@ InstructionCost BoUpSLP::getEntryCost(const TreeEntry *E,
       // initial vector or inserting a subvector.
       // TODO: Implement the analysis of the FirstInsert->getOperand(0)
       // subvector of ActualVecTy.
-      SmallVector<int> InsertMask(NumElts, UndefMaskElem);
-      copy(Mask, std::next(InsertMask.begin(), OffsetBeg));
       if (!isUndefVector(FirstInsert->getOperand(0), InsertMask) &&
           NumScalars != NumElts && !IsWholeSubvector) {
         if (InsertVecSz != VecSz) {

diff  --git a/llvm/test/Transforms/SLPVectorizer/X86/buildvector-insert-mask-size.ll b/llvm/test/Transforms/SLPVectorizer/X86/buildvector-insert-mask-size.ll
new file mode 100644
index 0000000000000..20bcbca61eb19
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/X86/buildvector-insert-mask-size.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+;RUN: opt -S -slp-vectorizer -mtriple=x86_64-unknown-linux < %s -slp-threshold=-1 | FileCheck %s
+
+define void @test() {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds float, ptr undef, i32 2
+; CHECK-NEXT:    [[TMP2:%.*]] = load <2 x float>, ptr [[TMP1]], align 4
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <2 x float> [[TMP2]], <2 x float> poison, <3 x i32> <i32 0, i32 undef, i32 1>
+; CHECK-NEXT:    store <3 x float> [[TMP3]], ptr null, align 4
+; CHECK-NEXT:    ret void
+;
+  %1 = getelementptr inbounds float, ptr undef, i32 2
+  %2 = load float, ptr %1, align 4
+  %3 = getelementptr inbounds float, ptr undef, i32 3
+  %4 = load float, ptr %3, align 4
+  %5 = insertelement <3 x float> poison, float %2, i64 0
+  %6 = insertelement <3 x float> %5, float %4, i64 2
+  store <3 x float> %6, ptr null, align 4
+  ret void
+}