[llvm] [LSROA] Support arrays, vectors, and nested structs (PR #218846)
Justin Bogner via llvm-commits
llvm-commits at lists.llvm.org
Wed Aug 26 00:01:11 PDT 2026
https://github.com/bogner created https://github.com/llvm/llvm-project/pull/218846
Update the algorithm to do a pre-order traversal of all accesses to a `structured.alloca` via `structured.gep` instead of just walking a single level of a struct.
There were some existing tests showing that we weren't handling a some of this and I also added a vector case since that falls out of the implementation.
Fixes #192619 and #192620
>From d495996baf1a504006afffe1ac841a13b7333bb3 Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail at justinbogner.com>
Date: Tue, 25 Aug 2026 23:50:33 -0700
Subject: [PATCH] [LSROA] Support arrays, vectors, and nested structs
Update the algorithm to do a pre-order traversal of all accesses to a
`structured.alloca` via `structured.gep` instead of just walking a
single level of a struct.
There were some existing tests showing that we weren't handling a some
of this and I also added a vector case since that falls out of the
implementation.
Fixes #192619 and #192620
---
llvm/include/llvm/IR/IntrinsicInst.h | 27 +++--
llvm/lib/Transforms/Scalar/LogicalSROA.cpp | 116 ++++++++++++++-------
llvm/test/Transforms/LSROA/array.ll | 7 +-
llvm/test/Transforms/LSROA/nesting.ll | 40 +++++--
llvm/test/Transforms/LSROA/vector.ll | 23 ++++
5 files changed, 149 insertions(+), 64 deletions(-)
create mode 100644 llvm/test/Transforms/LSROA/vector.ll
diff --git a/llvm/include/llvm/IR/IntrinsicInst.h b/llvm/include/llvm/IR/IntrinsicInst.h
index 7c0d9889a6d1a..262cd6c88a113 100644
--- a/llvm/include/llvm/IR/IntrinsicInst.h
+++ b/llvm/include/llvm/IR/IntrinsicInst.h
@@ -1809,23 +1809,20 @@ class StructuredGEPInst : public IntrinsicInst {
return make_range(op_begin() + 1, op_begin() + 1 + getNumIndices());
}
+ static Type *getTypeAtIndex(Type *Ty, const Value *Index) {
+ if (ArrayType *AT = dyn_cast<ArrayType>(Ty))
+ return AT->getElementType();
+ if (VectorType *VT = dyn_cast<VectorType>(Ty))
+ return VT->getElementType();
+ if (StructType *ST = dyn_cast<StructType>(Ty))
+ return ST->getElementType(cast<ConstantInt>(Index)->getZExtValue());
+ llvm_unreachable("llvm.structured.gep only indexes aggregate types");
+ }
+
Type *getResultElementType() const {
Type *CurrentType = getBaseType();
- for (unsigned I = 0; I < getNumIndices(); I++) {
- if (ArrayType *AT = dyn_cast<ArrayType>(CurrentType)) {
- CurrentType = AT->getElementType();
- } else if (VectorType *VT = dyn_cast<VectorType>(CurrentType)) {
- CurrentType = VT->getElementType();
- } else if (StructType *ST = dyn_cast<StructType>(CurrentType)) {
- ConstantInt *CI = cast<ConstantInt>(getIndexOperand(I));
- CurrentType = ST->getElementType(CI->getZExtValue());
- } else {
- // FIXME(Keenuts): add testing reaching those places once initial
- // implementation has landed.
- llvm_unreachable("unimplemented");
- }
- }
-
+ for (unsigned I = 0; I < getNumIndices(); I++)
+ CurrentType = getTypeAtIndex(CurrentType, getIndexOperand(I));
return CurrentType;
}
};
diff --git a/llvm/lib/Transforms/Scalar/LogicalSROA.cpp b/llvm/lib/Transforms/Scalar/LogicalSROA.cpp
index a0ed5cd0681d7..041b421d248f1 100644
--- a/llvm/lib/Transforms/Scalar/LogicalSROA.cpp
+++ b/llvm/lib/Transforms/Scalar/LogicalSROA.cpp
@@ -11,9 +11,6 @@
/// It tries to identify promotable elements of an aggregate alloca, and
/// promote them to multiple allocas of scalar type.
///
-/// FIXME: nested aggregates are not fully optimized (#192619).
-/// FIXME: array are not optimized (#192620).
-///
//===----------------------------------------------------------------------===//
#include "llvm/Transforms/Scalar/LogicalSROA.h"
@@ -113,27 +110,72 @@ static bool isAllocaSplittable(StructuredAllocaInst &SAI) {
return true;
}
-// Returns a vector with one element for each field of the struct allocated by
-// SAI. Each element is a vector of SGEP instruction referencing this field.
-// This function ignores lifetime intrinsics.
-static SmallVector<SmallVector<StructuredGEPInst *>>
-collectPerFieldSGEP(StructuredAllocaInst &SAI) {
- StructType *ST = cast<StructType>(SAI.getAllocationType());
- SmallVector<SmallVector<StructuredGEPInst *>> Output(ST->getNumElements());
-
- for (User *U : SAI.users()) {
- if (isa<LifetimeIntrinsic>(U))
+namespace {
+struct FieldSGEPs {
+ /// All GEPs that access this specific field.
+ SmallVector<StructuredGEPInst *, 1> SGEPs;
+ /// Type of the access.
+ Type *Ty;
+ /// The number of index arguments common to the collection of GEPs.
+ unsigned NumIndices;
+};
+} // namespace
+
+/// Returns a vector with one element for each field that is independently
+/// accessed of an SAI. Each element catalogues the list of GEPs for this field
+/// as well as the information needed to rewrite the GEP to a smaller alloca.
+/// This function ignores lifetime intrinsics.
+static SmallVector<FieldSGEPs> collectPerFieldSGEPs(StructuredAllocaInst &SAI) {
+ SmallVector<FieldSGEPs> PerFieldSGEPs;
+ SmallVector<FieldSGEPs> Worklist;
+
+ if (SAI.user_empty())
+ return PerFieldSGEPs;
+
+ Worklist.push_back({{}, SAI.getAllocationType(), /*NumSharedIndices=*/0});
+ for (User *U : SAI.users())
+ if (auto *SGEP = dyn_cast<StructuredGEPInst>(U))
+ Worklist.back().SGEPs.push_back(SGEP);
+
+ SmallVector<ConstantInt *> IndicesAtLevel;
+ while (!Worklist.empty()) {
+ FieldSGEPs Cur = Worklist.pop_back_val();
+
+ // When we run out of constant indices we're at the maximum depth we can
+ // split accesses at.
+ if (llvm::any_of(Cur.SGEPs, [&Cur](const auto *SGEP) {
+ return SGEP->getNumIndices() == Cur.NumIndices ||
+ !isa<ConstantInt>(SGEP->getIndexOperand(Cur.NumIndices));
+ })) {
+ PerFieldSGEPs.push_back(std::move(Cur));
continue;
+ }
- auto *SGEP = cast<StructuredGEPInst>(U);
-
- // IR rule: SGEP on struct can only use constant int as indices.
- ConstantInt *Index = cast<ConstantInt>(SGEP->getIndexOperand(0));
- assert(Index->getZExtValue() < Output.size());
- Output[Index->getZExtValue()].push_back(SGEP);
+ IndicesAtLevel.clear();
+ for (StructuredGEPInst *SGEP : Cur.SGEPs)
+ IndicesAtLevel.push_back(
+ cast<ConstantInt>(SGEP->getIndexOperand(Cur.NumIndices)));
+
+ // We need to operate on the unique indices that are accessed at this level
+ // of the GEPs. We sort by integer value rather than pointer identity so
+ // that the order we process these later will be deterministic.
+ llvm::sort(IndicesAtLevel, [](const auto &LHS, const auto &RHS) {
+ return LHS->getZExtValue() < RHS->getZExtValue();
+ });
+ IndicesAtLevel.erase(llvm::unique(IndicesAtLevel), IndicesAtLevel.end());
+
+ // Enqueue the next level of indices in pre-order.
+ for (const ConstantInt *CI : llvm::reverse(IndicesAtLevel)) {
+ Worklist.push_back({{},
+ StructuredGEPInst::getTypeAtIndex(Cur.Ty, CI),
+ Cur.NumIndices + 1});
+ for (StructuredGEPInst *SGEP : Cur.SGEPs)
+ if (SGEP->getIndexOperand(Cur.NumIndices) == CI)
+ Worklist.back().SGEPs.push_back(SGEP);
+ }
}
- return Output;
+ return PerFieldSGEPs;
}
// For each lifetime intrinsic in LifetimeIntrinsics, creates a new one, but
@@ -150,15 +192,16 @@ static void copyLifetimeIntrinsicFor(IRBuilder<> &B, LifetimeIntrinsic *II,
llvm_unreachable("invalid argument: expected a lifetime intrinsic");
}
-static void rewriteSGEPChain(IRBuilder<> &B, StructuredGEPInst *SGEP,
- StructuredAllocaInst *FieldAlloca) {
- if (SGEP->getNumIndices() == 1) {
+static void rewriteSGEPChain(IRBuilder<> &B, StructuredAllocaInst *FieldAlloca,
+ StructuredGEPInst *SGEP, unsigned NumIndices) {
+ if (SGEP->getNumIndices() == NumIndices) {
SGEP->replaceAllUsesWith(FieldAlloca);
SGEP->eraseFromParent();
return;
}
- SmallVector<Value *, 4> Indices(llvm::drop_begin(SGEP->indices()));
+ SmallVector<Value *, 4> Indices(
+ llvm::drop_begin(SGEP->indices(), NumIndices));
B.SetInsertPoint(SGEP);
auto *I = B.CreateStructuredGEP(FieldAlloca->getAllocationType(), FieldAlloca,
Indices, SGEP->getName());
@@ -167,32 +210,29 @@ static void rewriteSGEPChain(IRBuilder<> &B, StructuredGEPInst *SGEP,
}
static bool runOnStructuredAlloca(StructuredAllocaInst &SAI) {
- // For now, LogicalSROA only handles SGEP on structs.
- StructType *ST = dyn_cast<StructType>(SAI.getAllocationType());
- if (!ST)
+ Type *AllocaTy = SAI.getAllocationType();
+ // We only need to do anything with aggregate types.
+ if (!isa<ArrayType, StructType, VectorType>(AllocaTy))
return false;
if (!isAllocaSplittable(SAI))
return false;
- auto PerFieldSGEP = collectPerFieldSGEP(SAI);
- assert(PerFieldSGEP.size() == ST->getNumElements());
+ SmallVector<FieldSGEPs> PerFieldSGEPs = collectPerFieldSGEPs(SAI);
+ SmallVector<LifetimeIntrinsic *> LifetimeIntrinsics =
+ collectLifetimeIntrinsicsUsing(SAI);
- auto LifetimeIntrinsics = collectLifetimeIntrinsicsUsing(SAI);
IRBuilder B(&SAI);
- for (const auto &[FieldIndex, Users] : llvm::enumerate(PerFieldSGEP)) {
- if (Users.empty())
- continue;
-
+ for (const FieldSGEPs &Field : PerFieldSGEPs) {
B.SetInsertPoint(&SAI);
- auto *FieldAlloca = cast<StructuredAllocaInst>(
- B.CreateStructuredAlloca(ST->getElementType(FieldIndex)));
+ auto *FieldAlloca =
+ cast<StructuredAllocaInst>(B.CreateStructuredAlloca(Field.Ty));
for (auto II : LifetimeIntrinsics)
copyLifetimeIntrinsicFor(B, II, FieldAlloca);
- for (StructuredGEPInst *SGEP : Users)
- rewriteSGEPChain(B, SGEP, FieldAlloca);
+ for (StructuredGEPInst *SGEP : Field.SGEPs)
+ rewriteSGEPChain(B, FieldAlloca, SGEP, Field.NumIndices);
}
for (auto *II : LifetimeIntrinsics)
diff --git a/llvm/test/Transforms/LSROA/array.ll b/llvm/test/Transforms/LSROA/array.ll
index d488d90a38424..6e1c04cd3a139 100644
--- a/llvm/test/Transforms/LSROA/array.ll
+++ b/llvm/test/Transforms/LSROA/array.ll
@@ -9,10 +9,9 @@ declare ptr @llvm.structured.gep.p0(ptr, ...)
define i32 @test_simple_array() {
; CHECK-LABEL: define i32 @test_simple_array() {
; CHECK-NEXT: [[ENTRY:.*:]]
-; CHECK-NEXT: [[TMP:%.*]] = call elementtype([10 x i32]) ptr @llvm.structured.alloca.p0()
-; CHECK-NEXT: [[PTR:%.*]] = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype([10 x i32]) [[TMP]], i32 0)
-; CHECK-NEXT: store i32 0, ptr [[PTR]], align 4
-; CHECK-NEXT: [[RES:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT: [[TMP:%.*]] = call elementtype(i32) ptr @llvm.structured.alloca.p0()
+; CHECK-NEXT: store i32 0, ptr [[TMP]], align 4
+; CHECK-NEXT: [[RES:%.*]] = load i32, ptr [[TMP]], align 4
; CHECK-NEXT: ret i32 [[RES]]
;
entry:
diff --git a/llvm/test/Transforms/LSROA/nesting.ll b/llvm/test/Transforms/LSROA/nesting.ll
index 5291cfe4ebee0..dbeac4163778d 100644
--- a/llvm/test/Transforms/LSROA/nesting.ll
+++ b/llvm/test/Transforms/LSROA/nesting.ll
@@ -8,18 +8,15 @@ declare ptr @llvm.structured.gep.p0(ptr, ...)
%S = type { i32, { i32, i32 } }
-; TODO: It is fine to split the nested struct and discard the second element.
-; simply not implemented yet.
define i32 @test_nested_struct() {
; CHECK-LABEL: define i32 @test_nested_struct() {
; CHECK-NEXT: [[ENTRY:.*:]]
; CHECK-NEXT: [[TMP0:%.*]] = call elementtype(i32) ptr @llvm.structured.alloca.p0()
-; CHECK-NEXT: [[TMP1:%.*]] = call elementtype({ i32, i32 }) ptr @llvm.structured.alloca.p0()
-; CHECK-NEXT: [[TMP2:%.*]] = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype({ i32, i32 }) [[TMP1]], i32 0)
+; CHECK-NEXT: [[TMP1:%.*]] = call elementtype(i32) ptr @llvm.structured.alloca.p0()
; CHECK-NEXT: store i32 0, ptr [[TMP0]], align 4
-; CHECK-NEXT: store i32 1, ptr [[TMP2]], align 4
+; CHECK-NEXT: store i32 1, ptr [[TMP1]], align 4
; CHECK-NEXT: [[A:%.*]] = load i32, ptr [[TMP0]], align 4
-; CHECK-NEXT: [[B:%.*]] = load i32, ptr [[TMP2]], align 4
+; CHECK-NEXT: [[B:%.*]] = load i32, ptr [[TMP1]], align 4
; CHECK-NEXT: [[C:%.*]] = add i32 [[A]], [[B]]
; CHECK-NEXT: ret i32 [[C]]
;
@@ -28,7 +25,6 @@ entry:
%0 = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype(%S) %tmp, i32 0)
%1 = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype(%S) %tmp, i32 1, i32 0)
-
store i32 0, ptr %0
store i32 1, ptr %1
%a = load i32, ptr %0
@@ -37,3 +33,33 @@ entry:
%c = add i32 %a, %b
ret i32 %c
}
+
+%S2 = type { i32, [5 x %S] }
+
+define i32 @test_array_and_struct_nested() {
+; CHECK-LABEL: define i32 @test_array_and_struct_nested() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP1:%.*]] = call elementtype({ i32, i32 }) ptr @llvm.structured.alloca.p0()
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype({ i32, i32 }) [[TMP1]], i32 1)
+; CHECK-NEXT: [[P2:%.*]] = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype({ i32, i32 }) [[TMP1]], i32 0)
+; CHECK-NEXT: store i32 0, ptr [[TMP0]], align 4
+; CHECK-NEXT: store i32 1, ptr [[P2]], align 4
+; CHECK-NEXT: [[A:%.*]] = load i32, ptr [[TMP0]], align 4
+; CHECK-NEXT: [[B:%.*]] = load i32, ptr [[P2]], align 4
+; CHECK-NEXT: [[C:%.*]] = add i32 [[A]], [[B]]
+; CHECK-NEXT: ret i32 [[A]]
+;
+entry:
+ %tmp = call elementtype(%S2) ptr @llvm.structured.alloca.p0()
+ %p0 = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype(%S2) %tmp, i32 1, i32 1, i32 1, i32 1)
+ %p1 = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype(%S2) %tmp, i32 1, i32 1, i32 1)
+ %p2 = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype({ i32, i32 }) %p1, i32 0)
+
+ store i32 0, ptr %p0
+ store i32 1, ptr %p2
+ %a = load i32, ptr %p0
+ %b = load i32, ptr %p2
+
+ %c = add i32 %a, %b
+ ret i32 %a
+}
diff --git a/llvm/test/Transforms/LSROA/vector.ll b/llvm/test/Transforms/LSROA/vector.ll
new file mode 100644
index 0000000000000..6c9aad3695ad4
--- /dev/null
+++ b/llvm/test/Transforms/LSROA/vector.ll
@@ -0,0 +1,23 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes='logical-sroa' -S | FileCheck %s --check-prefixes=CHECK
+
+declare void @llvm.lifetime.start.p0(ptr nocapture)
+declare void @llvm.lifetime.end.p0(ptr nocapture)
+declare ptr @llvm.structured.alloca.p0()
+declare ptr @llvm.structured.gep.p0(ptr, ...)
+
+define i32 @test_simple_vector() {
+; CHECK-LABEL: define i32 @test_simple_vector() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP:%.*]] = call elementtype(i32) ptr @llvm.structured.alloca.p0()
+; CHECK-NEXT: store i32 0, ptr [[TMP]], align 4
+; CHECK-NEXT: [[RES:%.*]] = load i32, ptr [[TMP]], align 4
+; CHECK-NEXT: ret i32 [[RES]]
+;
+entry:
+ %tmp = call elementtype(<10 x i32>) ptr @llvm.structured.alloca.p0()
+ %ptr = call ptr (ptr, ...) @llvm.structured.gep.p0(ptr elementtype(<10 x i32>) %tmp, i32 0)
+ store i32 0, ptr %ptr
+ %res = load i32, ptr %ptr
+ ret i32 %res
+}
More information about the llvm-commits
mailing list