[llvm] Yonah/sroa integer aggregates late (PR #219353)

[Yonah Goldberg via llvm-commits]

https://github.com/YonahGoldberg created https://github.com/llvm/llvm-project/pull/219353

None

>From 28824a57bde6b30b4ac20e1a8e065a7b9f1e7c9f Mon Sep 17 00:00:00 2001
From: Yonah Goldberg <ygoldberg at nvidia.com>
Date: Thu, 27 Aug 2026 07:23:46 +0000
Subject: [PATCH 1/2] [SROA] Prefer legal integers for aggregate partitions

---
 llvm/lib/Transforms/Scalar/SROA.cpp | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index e4770154e2998..a449b1c9ff61e 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -4804,6 +4804,14 @@ static Type *getTypePartition(const DataLayout &DL, Type *Ty, uint64_t Offset,
   return SubTy;
 }
 
+static bool containsNonIntegralPointer(Type *Ty, const DataLayout &DL) {
+  if (auto *PtrTy = dyn_cast<PointerType>(Ty))
+    return DL.isNonIntegralPointerType(PtrTy);
+  return llvm::any_of(Ty->subtypes(), [&](Type *SubTy) {
+    return containsNonIntegralPointer(SubTy, DL);
+  });
+}
+
 /// Pre-split loads and stores to simplify rewriting.
 ///
 /// We want to break up the splittable load+store pairs as much as
@@ -5458,11 +5466,11 @@ selectPartitionType(Partition &P, const DataLayout &DL, AllocaInst &AI,
   // type?
   if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
                                                P.beginOffset(), P.size())) {
-    // If the partition is an integer array that can be spanned by a legal
-    // integer type, prefer to represent it as a legal integer type because
-    // it's more likely to be promotable.
-    if (TypePartitionTy->isArrayTy() &&
-        TypePartitionTy->getArrayElementType()->isIntegerTy() &&
+    // If the partition is an aggregate without non-integral pointers that can
+    // be spanned by a legal integer type, prefer to represent it as a legal
+    // integer type because it's more likely to be promotable.
+    if (TypePartitionTy->isAggregateType() &&
+        !containsNonIntegralPointer(TypePartitionTy, DL) &&
         DL.isLegalInteger(P.size() * 8))
       TypePartitionTy = Type::getIntNTy(C, P.size() * 8);
     // There was no common type used, so we prefer integer widening promotion.

>From 4380f8a332adccbc9a92385d8a84a2885e00585f Mon Sep 17 00:00:00 2001
From: Yonah Goldberg <ygoldberg at nvidia.com>
Date: Fri, 28 Aug 2026 03:27:33 +0000
Subject: [PATCH 2/2] [SROA] Canonicalize aggregate partitions to integers late

---
 llvm/lib/Transforms/Scalar/SROA.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index a449b1c9ff61e..14f029f36a236 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -5466,10 +5466,10 @@ selectPartitionType(Partition &P, const DataLayout &DL, AllocaInst &AI,
   // type?
   if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
                                                P.beginOffset(), P.size())) {
-    // If the partition is an aggregate without non-integral pointers that can
-    // be spanned by a legal integer type, prefer to represent it as a legal
-    // integer type because it's more likely to be promotable.
-    if (TypePartitionTy->isAggregateType() &&
+    // If aggregate canonicalization is requested and the partition can be
+    // spanned by a legal integer type, prefer the integer type because it's
+    // more likely to be promotable.
+    if (AggregateToVector && TypePartitionTy->isAggregateType() &&
         !containsNonIntegralPointer(TypePartitionTy, DL) &&
         DL.isLegalInteger(P.size() * 8))
       TypePartitionTy = Type::getIntNTy(C, P.size() * 8);