[flang-commits] [flang] [flang] Preserve the alignment attribute across stack/heap conversions (PR #234898)
Eugene Epshteyn via flang-commits
flang-commits at lists.llvm.org
Mon Oct 26 06:46:53 PDT 2026
https://github.com/eugeneepshteyn updated https://github.com/llvm/llvm-project/pull/234898
>From 359ddc711f8b2866a22e7d03308add0cea4f3fc5 Mon Sep 17 00:00:00 2001
From: Eugene Epshteyn <eepshteyn at nvidia.com>
Date: Thu, 22 Oct 2026 19:27:49 -0700
Subject: [PATCH 1/2] [flang] Preserve the alignment attribute across
stack/heap conversions
The two helpers that convert allocations between the stack and the heap
rebuilt the new allocation from the type, names and operands only, so a
requested alignment was lost in both directions:
- `fir::createAllocMemFromAlloca` (used by MemoryAllocation,
AllocationPlacement and the CUDA heap promotion) now copies the
`fir.alloca` alignment onto the `fir.allocmem` it creates;
- `AllocMemConversion::insertAlloca` (StackArrays, also reused by
AllocationPlacement) now copies the `fir.allocmem` alignment onto the
`fir.alloca` it creates.
AllocationPlacement's stack budget also accounts for the worst-case
padding an over-aligned stack object may need.
Today the only producer of the `fir.allocmem` alignment attribute is
the allocatable/pointer allocation path, which StackArrays does not
touch by design, so this fixes a latent rather than an observable drop;
it becomes live as soon as temporaries and user variables carry an
alignment.
---
.../Transforms/AllocationPlacement.cpp | 10 ++++++--
.../lib/Optimizer/Transforms/MemoryUtils.cpp | 12 ++++++----
.../lib/Optimizer/Transforms/StackArrays.cpp | 3 +++
flang/test/Fir/memory-allocation-opt.fir | 14 +++++++++++
.../test/Transforms/allocation-placement.fir | 24 +++++++++++++++++++
flang/test/Transforms/stack-arrays.fir | 15 ++++++++++++
6 files changed, 72 insertions(+), 6 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index 89783d78148316..a59c7181a1fef4 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -215,9 +215,15 @@ void AllocationPlacementPass::runOnOperation() {
? placementHook(info, basePolicy, stackBytesUsed)
: fir::decideAllocationPlacement(info, basePolicy, stackBytesUsed);
- // Account for the decision in the running stack budget.
- if (endsUpOnStack(placement, info.isCurrentlyOnStack) && info.byteSize)
+ // Account for the decision in the running stack budget, including the
+ // worst-case padding a requested alignment may need.
+ if (endsUpOnStack(placement, info.isCurrentlyOnStack) && info.byteSize) {
stackBytesUsed += static_cast<std::size_t>(*info.byteSize);
+ std::optional<std::int64_t> alignment =
+ alloca ? alloca.getAlignment() : allocmem.getAlignment();
+ if (alignment && *alignment > 1)
+ stackBytesUsed += static_cast<std::size_t>(*alignment - 1);
+ }
LLVM_DEBUG({
llvm::dbgs() << "allocation-placement: "
diff --git a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
index 00a25a7326f2dd..6da0c3141b8c44 100644
--- a/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
+++ b/flang/lib/Optimizer/Transforms/MemoryUtils.cpp
@@ -322,10 +322,14 @@ fir::AllocMemOp fir::createAllocMemFromAlloca(mlir::OpBuilder &builder,
return *opt;
return {};
};
- return fir::AllocMemOp::create(builder, alloca.getLoc(), alloca.getInType(),
- unpackName(alloca.getUniqName()),
- unpackName(alloca.getBindcName()),
- alloca.getTypeparams(), alloca.getShape());
+ auto heap = fir::AllocMemOp::create(
+ builder, alloca.getLoc(), alloca.getInType(),
+ unpackName(alloca.getUniqName()), unpackName(alloca.getBindcName()),
+ alloca.getTypeparams(), alloca.getShape());
+ // Keep a requested alignment across the stack-to-heap conversion.
+ if (std::optional<std::int64_t> alignment = alloca.getAlignment())
+ heap.setAlignment(*alignment);
+ return heap;
}
/// Device code keeps its stack allocations: the unified/managed entry points
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 0674cc24319606..6b8edde0c2d927 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -711,6 +711,9 @@ fir::AllocMemConversion::insertAlloca(fir::AllocMemOp &oldAlloc,
auto alloca =
fir::AllocaOp::create(rewriter, loc, varTy, uniqName, bindcName,
oldAlloc.getTypeparams(), oldAlloc.getShape());
+ // Keep a requested alignment across the heap-to-stack conversion.
+ if (std::optional<std::int64_t> alignment = oldAlloc.getAlignment())
+ alloca.setAlignment(*alignment);
if (emitLifetimeMarkers)
insertLifetimeMarkers(oldAlloc, alloca, rewriter);
diff --git a/flang/test/Fir/memory-allocation-opt.fir b/flang/test/Fir/memory-allocation-opt.fir
index 610a3f0b3072a6..9993ebc6b7e502 100644
--- a/flang/test/Fir/memory-allocation-opt.fir
+++ b/flang/test/Fir/memory-allocation-opt.fir
@@ -32,3 +32,17 @@ func.func @_QPs2(%arg0: !fir.ref<i32>) {
}
func.func private @_QPs3(!fir.ref<!fir.array<1000123xi32>>)
+// A requested alignment survives the move to the heap.
+
+// CHECK-LABEL: func @_QPs4(
+// CHECK: %[[mem:.*]] = fir.allocmem !fir.array<1000123xi32> <{alignment = 64 : i64, bindc_name = "array", uniq_name = "_QFs4Earray"}>
+// CHECK: fir.call @_QPs3(
+// CHECK: fir.freemem %[[mem]]
+// CHECK-NEXT: return
+
+func.func @_QPs4() {
+ %0 = fir.alloca !fir.array<1000123xi32> {alignment = 64 : i64, bindc_name = "array", uniq_name = "_QFs4Earray"}
+ fir.call @_QPs3(%0) : (!fir.ref<!fir.array<1000123xi32>>) -> ()
+ return
+}
+
diff --git a/flang/test/Transforms/allocation-placement.fir b/flang/test/Transforms/allocation-placement.fir
index aa09d830fd5a5f..c9c8f6d65e6374 100644
--- a/flang/test/Transforms/allocation-placement.fir
+++ b/flang/test/Transforms/allocation-placement.fir
@@ -63,6 +63,30 @@ func.func @big_temp_alloca() {
return
}
+// A requested alignment survives heap-to-stack ...
+// CHECK-LABEL: func.func @small_temp_aligned
+// CHECK: fir.alloca !fir.array<10xi32> <{alignment = 64 : i64}>
+// CHECK-NOT: fir.allocmem
+func.func @small_temp_aligned() {
+ %0 = fir.allocmem !fir.array<10xi32> {alignment = 64 : i64}
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<10xi32>>) -> !fir.ref<!fir.array<10xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<10xi32>>
+ return
+}
+
+// ... and stack-to-heap.
+// CHECK-LABEL: func.func @big_temp_alloca_aligned
+// CHECK: fir.allocmem !fir.array<100xi32> <{alignment = 64 : i64
+// CHECK: fir.freemem
+func.func @big_temp_alloca_aligned() {
+ %0 = fir.alloca !fir.array<100xi32> {alignment = 64 : i64}
+ return
+}
+
// Runtime-sized user variable (automatic array) -> heap.
// CHECK-LABEL: func.func @dyn_user
// CHECK: fir.allocmem !fir.array<?xi32>
diff --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index 5b149bfd8b6216..08d893d5679ea3 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -15,6 +15,21 @@ func.func @simple() {
// CHECK: fir.alloca !fir.array<42xi32>
// CHECK: return
+// A requested alignment survives the move to the stack
+func.func @aligned() {
+ %0 = fir.allocmem !fir.array<42xi32> {alignment = 64 : i64}
+ %c0_a = arith.constant 0 : index
+ %c0_i32_a = arith.constant 0 : i32
+ %ref_a = fir.convert %0 : (!fir.heap<!fir.array<42xi32>>) -> !fir.ref<!fir.array<42xi32>>
+ %elt_a = fir.coordinate_of %ref_a, %c0_a : (!fir.ref<!fir.array<42xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32_a to %elt_a : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<42xi32>>
+ return
+}
+// CHECK-LABEL: func.func @aligned()
+// CHECK: fir.alloca !fir.array<42xi32> <{alignment = 64 : i64}>
+// CHECK: return
+
// Check fir.must_be_heap allocations are not moved
func.func @must_be_heap() {
%0 = fir.allocmem !fir.array<42xi32> {fir.must_be_heap = true}
>From 74463d48b1f47af6253f564d8d121b6dca2e0a25 Mon Sep 17 00:00:00 2001
From: Eugene Epshteyn <eepshteyn at nvidia.com>
Date: Mon, 26 Oct 2026 04:31:29 -0700
Subject: [PATCH 2/2] [flang] Charge alignment padding in the placement
decision, not only after it
The worst-case padding a requested alignment may need (alignment - 1)
was added to the running stack budget after the placement decision, so
it affected the allocations that followed but not whether the aligned
allocation itself fit the per-function limit.
Record the requested alignment in PendingAllocationInfo and compute the
stack bytes an allocation needs in one place, stackBytesNeeded(), used
both by shouldAllocateOnStack for the total-stack-limit comparison and
by the placement pass when charging the budget. The small-array
threshold still compares the object's own size. Tests cover an aligned
array that no longer fits because of its padding, and the padding of a
placed array being charged against the next one.
---
.../Optimizer/Support/AllocationPolicy.h | 8 ++++
.../Optimizer/Support/AllocationPolicy.cpp | 17 +++++++-
.../Transforms/AllocationPlacement.cpp | 15 +++----
.../allocation-placement-budget.fir | 40 +++++++++++++++++++
4 files changed, 69 insertions(+), 11 deletions(-)
diff --git a/flang/include/flang/Optimizer/Support/AllocationPolicy.h b/flang/include/flang/Optimizer/Support/AllocationPolicy.h
index 0b9a3be1698755..b49733afda5095 100644
--- a/flang/include/flang/Optimizer/Support/AllocationPolicy.h
+++ b/flang/include/flang/Optimizer/Support/AllocationPolicy.h
@@ -84,12 +84,20 @@ struct PendingAllocationInfo {
bool isDynamic = false;
/// The constant size of the allocation in bytes, if it can be determined.
std::optional<std::int64_t> byteSize;
+ /// The alignment requested for the allocation in bytes, if any. The stack
+ /// budget charges the worst-case padding (alignment - 1) on top of byteSize.
+ std::optional<std::int64_t> alignment;
/// An operation at the point where the allocation lives or will be inserted,
/// used for context dependent decisions such as code that runs on a device.
/// Optional: without it the offload region rule is not applied.
mlir::Operation *context = nullptr;
};
+/// The stack bytes the allocation described by \p info would consume, if its
+/// size is known: its byteSize plus the worst-case padding for its requested
+/// alignment. The same figure decides placement and is charged to the budget.
+std::optional<std::size_t> stackBytesNeeded(const PendingAllocationInfo &info);
+
/// Module-level information needed to compute constant allocation sizes.
struct AllocationSizeContext {
std::optional<mlir::DataLayout> dataLayout;
diff --git a/flang/lib/Optimizer/Support/AllocationPolicy.cpp b/flang/lib/Optimizer/Support/AllocationPolicy.cpp
index 1e460d18096d91..2fd913bd4524af 100644
--- a/flang/lib/Optimizer/Support/AllocationPolicy.cpp
+++ b/flang/lib/Optimizer/Support/AllocationPolicy.cpp
@@ -64,6 +64,16 @@ bool fir::shouldUseStackForCopyin(mlir::Location loc, mlir::Type sequenceType,
return shouldAllocateOnStack(info, copyInPolicy, /*stackBytesUsed=*/0);
}
+std::optional<std::size_t>
+fir::stackBytesNeeded(const PendingAllocationInfo &info) {
+ if (!info.byteSize)
+ return std::nullopt;
+ auto bytes = static_cast<std::size_t>(*info.byteSize);
+ if (info.alignment && *info.alignment > 1)
+ bytes += static_cast<std::size_t>(*info.alignment - 1);
+ return bytes;
+}
+
/// The policy in effect for \p info. Inside an offload region -fstack-arrays is
/// not honored, since the device stack is far smaller than the host one.
static fir::AllocationPolicy policyFor(const fir::PendingAllocationInfo &info,
@@ -100,9 +110,12 @@ static bool shouldAllocateOnStackImpl(const fir::PendingAllocationInfo &info,
auto size = static_cast<std::size_t>(*info.byteSize);
// Small array temporaries go on the stack while the per-function budget
- // allows it; bigger ones go on the heap.
+ // allows it; bigger ones go on the heap. The size threshold looks at the
+ // object itself; the budget also charges the worst-case padding a requested
+ // alignment may need.
return size <= policy.smallArrayThresholdBytes &&
- stackBytesUsed + size <= policy.totalStackLimitBytes;
+ stackBytesUsed + *stackBytesNeeded(info) <=
+ policy.totalStackLimitBytes;
}
bool fir::shouldAllocateOnStack(const PendingAllocationInfo &info,
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index a59c7181a1fef4..40605ffd3cc31d 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -207,6 +207,7 @@ void AllocationPlacementPass::runOnOperation() {
alloca ? alloca.isDynamic()
: (allocmem.hasLenParams() || allocmem.hasShapeOperands());
info.byteSize = getConstantByteSize(op, dl, kindMap);
+ info.alignment = alloca ? alloca.getAlignment() : allocmem.getAlignment();
// A hook, if provided, fully overrides the default policy; it may delegate
// back to decideAllocationPlacement after adjusting the policy.
@@ -215,15 +216,11 @@ void AllocationPlacementPass::runOnOperation() {
? placementHook(info, basePolicy, stackBytesUsed)
: fir::decideAllocationPlacement(info, basePolicy, stackBytesUsed);
- // Account for the decision in the running stack budget, including the
- // worst-case padding a requested alignment may need.
- if (endsUpOnStack(placement, info.isCurrentlyOnStack) && info.byteSize) {
- stackBytesUsed += static_cast<std::size_t>(*info.byteSize);
- std::optional<std::int64_t> alignment =
- alloca ? alloca.getAlignment() : allocmem.getAlignment();
- if (alignment && *alignment > 1)
- stackBytesUsed += static_cast<std::size_t>(*alignment - 1);
- }
+ // Account for the decision in the running stack budget, with the same
+ // figure the decision used (size plus worst-case alignment padding).
+ if (endsUpOnStack(placement, info.isCurrentlyOnStack))
+ if (std::optional<std::size_t> bytes = fir::stackBytesNeeded(info))
+ stackBytesUsed += *bytes;
LLVM_DEBUG({
llvm::dbgs() << "allocation-placement: "
diff --git a/flang/test/Transforms/allocation-placement-budget.fir b/flang/test/Transforms/allocation-placement-budget.fir
index 6da58c4cfe9be1..1d7aafbe452e00 100644
--- a/flang/test/Transforms/allocation-placement-budget.fir
+++ b/flang/test/Transforms/allocation-placement-budget.fir
@@ -26,4 +26,44 @@ func.func @budget() {
return
}
+// A requested alignment counts against the budget as worst-case padding
+// (alignment - 1), both when deciding and when accounting: the 40-byte array
+// is below the small-array threshold, but 40 + 63 bytes of padding do not fit
+// the 64-byte limit, so it stays on the heap.
+// CHECK-LABEL: func.func @budget_aligned_too_big
+// CHECK: fir.allocmem !fir.array<10xi32> <{alignment = 64 : i64
+// CHECK-NOT: fir.alloca
+func.func @budget_aligned_too_big() {
+ %0 = fir.allocmem !fir.array<10xi32> {alignment = 64 : i64}
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r0 = fir.convert %0 : (!fir.heap<!fir.array<10xi32>>) -> !fir.ref<!fir.array<10xi32>>
+ %e0 = fir.coordinate_of %r0, %c0 : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e0 : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<10xi32>>
+ return
+}
+
+// An 8-byte array aligned to 32 fits (8 + 31 = 39 <= 64) and goes on the
+// stack; its padding is then charged, so the following 40-byte array no longer
+// fits (39 + 40 = 79 > 64) and stays on the heap.
+// CHECK-LABEL: func.func @budget_aligned_padding_charged
+// CHECK-DAG: fir.alloca !fir.array<2xi32> <{alignment = 32 : i64}>
+// CHECK-DAG: fir.allocmem !fir.array<10xi32>
+func.func @budget_aligned_padding_charged() {
+ %0 = fir.allocmem !fir.array<2xi32> {alignment = 32 : i64}
+ %1 = fir.allocmem !fir.array<10xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r0 = fir.convert %0 : (!fir.heap<!fir.array<2xi32>>) -> !fir.ref<!fir.array<2xi32>>
+ %e0 = fir.coordinate_of %r0, %c0 : (!fir.ref<!fir.array<2xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e0 : !fir.ref<i32>
+ %r1 = fir.convert %1 : (!fir.heap<!fir.array<10xi32>>) -> !fir.ref<!fir.array<10xi32>>
+ %e1 = fir.coordinate_of %r1, %c0 : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e1 : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<2xi32>>
+ fir.freemem %1 : !fir.heap<!fir.array<10xi32>>
+ return
+}
+
}
More information about the flang-commits
mailing list