[llvm] 8faeed0 - [SROA] Add Stored Value Size Check for Tree-Structured Merge (#162921)

[via llvm-commits]

Author: Chengjun
Date: 2025-10-11T00:23:04Z
New Revision: 8faeed042af2c1278cc71f0a5ef1a4d45f905eb8

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

LOG: [SROA] Add Stored Value Size Check for Tree-Structured Merge (#162921)

The change fixes a bug in the SROA where tree-structured merge
optimization was incorrectly applied when the size of the stored value
was not a multiple of the new allocated element type size. The original
change is https://github.com/llvm/llvm-project/pull/152793. A simple
repro would be

```
define <1 x i32> @foo(<1 x i16> %a, <1 x i16> %b) {
entry:
  %alloca = alloca [1 x i32]

  %ptr0 = getelementptr inbounds [2 x i16], ptr %alloca, i32 0, i32 0
  store <1 x i16> %a, ptr %ptr0

  %ptr1 = getelementptr inbounds [2 x i16], ptr %alloca, i32 0, i32 1
  store <1 x i16> %b, ptr %ptr1

  %result = load <1 x i32>, ptr %alloca
  ret <1 x i32> %result
}
```

Currently, this will lead to a compile time crash.

In this change, we will skip the tree-structured merge for this case and
fall back to normal SROA.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/SROA.cpp
    llvm/test/Transforms/SROA/vector-promotion-cannot-tree-structure-merge.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 45d3d493a9e68..b9d332b590368 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2961,6 +2961,7 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
         isa<FixedVectorType>(NewAI.getAllocatedType())
             ? cast<FixedVectorType>(NewAI.getAllocatedType())->getElementType()
             : Type::getInt8Ty(NewAI.getContext());
+    unsigned AllocatedEltTySize = DL.getTypeSizeInBits(AllocatedEltTy);
 
     // Helper to check if a type is
     //  1. A fixed vector type
@@ -2991,10 +2992,17 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
         // Do not handle the case if
         //   1. The store does not meet the conditions in the helper function
         //   2. The store is volatile
+        //   3. The total store size is not a multiple of the allocated element
+        //   type size
         if (!IsTypeValidForTreeStructuredMerge(
                 SI->getValueOperand()->getType()) ||
             SI->isVolatile())
           return std::nullopt;
+        auto *VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
+        unsigned NumElts = VecTy->getNumElements();
+        unsigned EltSize = DL.getTypeSizeInBits(VecTy->getElementType());
+        if (NumElts * EltSize % AllocatedEltTySize != 0)
+          return std::nullopt;
         StoreInfos.emplace_back(SI, S.beginOffset(), S.endOffset(),
                                 SI->getValueOperand());
       } else {

diff  --git a/llvm/test/Transforms/SROA/vector-promotion-cannot-tree-structure-merge.ll b/llvm/test/Transforms/SROA/vector-promotion-cannot-tree-structure-merge.ll
index c858d071451e8..ead6e027ed37c 100644
--- a/llvm/test/Transforms/SROA/vector-promotion-cannot-tree-structure-merge.ll
+++ b/llvm/test/Transforms/SROA/vector-promotion-cannot-tree-structure-merge.ll
@@ -219,4 +219,18 @@ entry:
 
 }
 
+define <1 x i32> @test_store_value_size_not_multiple_of_allocated_element_type_size(<1 x i16> %a, <1 x i16> %b) {
+entry:
+  %alloca = alloca [2 x i16]
+
+  %ptr0 = getelementptr inbounds [2 x i16], ptr %alloca, i32 0, i32 0
+  store <1 x i16> %a, ptr %ptr0
+
+  %ptr1 = getelementptr inbounds [2 x i16], ptr %alloca, i32 0, i32 1
+  store <1 x i16> %b, ptr %ptr1
+
+  %result = load <1 x i32>, ptr %alloca
+  ret <1 x i32> %result
+}
+
 declare void @llvm.memset.p0.i64(ptr nocapture writeonly, i8, i64, i1 immarg)