[SROA] Refactor rewritePartition alloca type selection process

[YonahGoldberg]

When an aggregate alloca has mixed scalar type uses (e.g., store i32 and load float), SROA needs to find a common type for the partition.
SROA's type selection logic works as follows:

1. First, try to find a common type used by all uses
2. If that fails, try to extract an appropriate subtype from the original allocated type via getTypePartition
3. If still no type, try to use the largest bitwidth integer type used
4. If that fails etc... etc...

The problem occurs in step 2: getTypePartition can return an aggregate type (like [2 x half]) that spans the partition, but aggregate types are not promotable to SSA values (they are not single-value types). This prevents SROA from eliminating the alloca.

For example, given:

%alloca = alloca [2 x half]
store i32 42, ptr %alloca
%val = load float, ptr %alloca

Previously, SROA would:

• Find no common type between i32 and float
• Use getTypePartition which returns [2 x half]
• Create a new alloca [2 x half]
• Fail to promote it because arrays are not single-value types

This PR modifies step 3 to also trigger when getTypePartition returns a non-promotable aggregate (checked via !isSingleValueType()). This causes SROA to prefer the integer type (i32) over the aggregate type, allowing the alloca to be fully promoted and eliminated.

After this change, the example is optimized to:

%0 = bitcast i32 42 to float
ret void

The alloca is completely eliminated.

This PR also allows [2 x float] allocas to get treated as a i64 partition, which allows promotion. I guess the point of not allowing this before was that you'd have to have an extra bitcast potentially to float, but with opaque pointers this isn't true anymore. This resolves #164308.

I read through the discussion on the commit that implemented some of this logic and I'm wondering if there's a larger refactoring that should happen. For example this comment:

Agreed. But until LLVM removes pointer sub-types it's convenient to get the alloca type right to avoid bitcast on every access anyway.
When pointer sub-types goes away, I guess all this code in SROA to find the right type for alloca would go away, but as you say it would have to be replaced with code to get the right load/store type instead. (FWIW Alive2's alloca only takes the number of bytes to allocate as argument)
So I see this patch as a step in the right direction.

Is some of the code that exists here not useful anymore now that we have opaque pointers? I guess it doesn't matter as much what the type of the alloca is.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yonah Goldberg (YonahGoldberg)

Changes

When an aggregate alloca has mixed scalar type uses (e.g., store i32 and load float), SROA needs to find a common type for the partition.
SROA's type selection logic works as follows:

1. First, try to find a common type used by all uses
2. If that fails, try to extract an appropriate subtype from the original allocated type via getTypePartition
3. If still no type, try to use the largest bitwidth integer type used
4. If that fails etc... etc...

The problem occurs in step 2: getTypePartition can return an aggregate type (like [2 x half]) that spans the partition, but aggregate types are not promotable to SSA values (they are not single-value types). This prevents SROA from eliminating the alloca.

For example, given:

%alloca = alloca [2 x half]
store i32 42, ptr %alloca
%val = load float, ptr %alloca

Previously, SROA would:

• Find no common type between i32 and float
• Use getTypePartition which returns [2 x half]
• Create a new alloca [2 x half]
• Fail to promote it because arrays are not single-value types

This PR modifies step 3 to also trigger when getTypePartition returns a non-promotable aggregate (checked via !isSingleValueType()). This causes SROA to prefer the integer type (i32) over the aggregate type, allowing the alloca to be fully promoted and eliminated.

After this change, the example is optimized to:

%0 = bitcast i32 42 to float
ret void

The alloca is completely eliminated.

I read through the discussion on the commit that implemented some of this logic and I'm wondering if there's a larger refactoring that should happen. For example this comment:

> Agreed. But until LLVM removes pointer sub-types it's convenient to get the alloca type right to avoid bitcast on every access anyway.
When pointer sub-types goes away, I guess all this code in SROA to find the right type for alloca would go away, but as you say it would have to be replaced with code to get the right load/store type instead. (FWIW Alive2's alloca only takes the number of bytes to allocate as argument)
So I see this patch as a step in the right direction.

Is some of the code that exists here not useful anymore now that we have opaque pointers? I guess it doesn't matter as much what the type of the alloca is.

---

Full diff: https://github.com/llvm/llvm-project/pull/167771.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/SROA.cpp (+3-1)
• (added) llvm/test/Transforms/SROA/prefer-integer-partition.ll (+17)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 5c60fad6f91aa..7905cfe95336d 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -5234,7 +5234,9 @@ AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
       SliceTy = TypePartitionTy;
 
   // If still not, can we use the largest bitwidth integer type used?
-  if (!SliceTy && CommonUseTy.second)
+  // If SliceTy is a non-promotable aggregate, prefer to represent as an integer type
+  // because it's more likely to be promotable.
+  if ((!SliceTy || !SliceTy->isSingleValueType()) && CommonUseTy.second)
     if (DL.getTypeAllocSize(CommonUseTy.second).getFixedValue() >= P.size()) {
       SliceTy = CommonUseTy.second;
       SliceVecTy = dyn_cast<VectorType>(SliceTy);
diff --git a/llvm/test/Transforms/SROA/prefer-integer-partition.ll b/llvm/test/Transforms/SROA/prefer-integer-partition.ll
new file mode 100644
index 0000000000000..3606af8debd69
--- /dev/null
+++ b/llvm/test/Transforms/SROA/prefer-integer-partition.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=sroa -S | FileCheck %s
+
+; Ensure that the [2 x half] alloca is spanned by an i32 partition.
+
+define void @test() {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32 42 to float
+; CHECK-NEXT:    ret void
+;
+entry:
+  %alloca = alloca [2 x half]
+  store i32 42, ptr %alloca
+  %val = load float, ptr %alloca
+  ret void
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[YonahGoldberg]

This is related to issue: #164308, but it doesn't seem like it fixes it, so I will investigate this one as well.

[YonahGoldberg]

Updated the PR to fix the Julia issue as well.

[YonahGoldberg]

FYI I am looking into a larger refactor to simplify the type selection process in rewritePartition

[github-actions]

🐧 Linux x64 Test Results

• 193491 tests passed
• 6234 tests skipped

✅ The build succeeded and all tests passed.

[YonahGoldberg]

Both of the arm changes are due to the early exit I added if isVectorPromotionViable returns a vector of floats. The neon testcase has the following pattern:

alloca i64
load <4 x bfloat>
store i64

Originally we prefer to promote via integer widening, now we prefer to promote via vector promotion. I think this is a win because if you're loading as a float type its probably being treated as a float.

The bf16 testcase has a alloca <1 x double> that we now promote via vector promotion instead of integer widening, which causes the small diff. Again, I think this is fine.