[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
Fri Oct 23 03:25:38 PDT 2026


https://github.com/eugeneepshteyn created https://github.com/llvm/llvm-project/pull/234898

Three places move an allocation between the stack and the heap and drop any `alignment` attribute on the way: `createAllocMemFromAlloca` in MemoryUtils, `insertAlloca` in StackArrays, and the stack-budget accounting in AllocationPlacement. This patch copies the attribute across in the first two and lets the third budget for the padding an over-aligned alloca may need.

Nothing user-visible depends on this today: the only producer of an aligned `fir.allocmem` is an ALLOCATE statement, which StackArrays excludes by design, and flang's own temporaries carry no alignment. It is groundwork for the alignment-directive series, where aligned locals and temporaries will flow through these passes.

Stacked on #234897.


>From c86795ec5ba55102235a4a28876baa7cb5955a4e Mon Sep 17 00:00:00 2001
From: Eugene Epshteyn <eepshteyn at nvidia.com>
Date: Thu, 22 Oct 2026 19:27:48 -0700
Subject: [PATCH 1/2] [flang] Add an alignment attribute to fir.alloca

`fir.allocmem` and `fir.global` can carry a requested alignment, but
`fir.alloca` could not, so there was no way to represent an aligned
stack object in FIR. Add an optional `alignment` attribute to
`fir.alloca`, verified to be a positive power of two like the one on
`fir.allocmem`, and forward it to `llvm.alloca` in codegen.

The requested value is a minimum: codegen raises it to the ABI alignment
of the allocated type when it is smaller, so an alloca never ends up
less aligned than one without the attribute. The same clamp is applied
to `fir.global`, which until now forwarded the attribute verbatim.

This is the first step towards a directive that sets the alignment of
local, automatic and global variables; passes that convert between
stack and heap allocations learn to preserve the attribute in the
following change.
---
 .../include/flang/Optimizer/Dialect/FIROps.td |  8 ++++-
 flang/lib/Optimizer/CodeGen/CodeGen.cpp       | 15 +++++++--
 flang/lib/Optimizer/Dialect/FIROps.cpp        | 16 ++++++----
 flang/test/Fir/convert-to-llvm.fir            | 32 +++++++++++++++++++
 flang/test/Fir/fir-ops.fir                    |  2 ++
 flang/test/Fir/invalid.fir                    |  8 +++++
 6 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index fc00daac1695b1..0b727b46dc4ce4 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -205,6 +205,11 @@ def fir_AllocaOp : fir_Op<"alloca", [
 
     The pinned attribute is used to flag fir.alloca operation in a specific
     region and avoid them being hoisted in an alloca hoisting pass.
+
+    The optional alignment attribute requests a minimum byte alignment for
+    the allocation, as `fir.allocmem` and `fir.global` do. It must be a
+    positive power of two. Code generation never lowers the alignment below
+    the ABI alignment of the allocated type.
   }];
 
   let arguments = (ins
@@ -213,7 +218,8 @@ def fir_AllocaOp : fir_Op<"alloca", [
     OptionalAttr<StrAttr>:$bindc_name,
     UnitAttr:$pinned,
     Variadic<AnyIntegerType>:$typeparams,
-    Variadic<AnyIntegerType>:$shape
+    Variadic<AnyIntegerType>:$shape,
+    OptionalAttr<I64Attr>:$alignment
   );
 
   let results =
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index d5097aa6f76035..63cc61193d1f65 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -480,13 +480,21 @@ struct AllocaOpConversion : public fir::FIROpConversion<fir::AllocaOp> {
       rewriter.setInsertionPointAfter(size.getDefiningOp());
     }
 
+    // A requested alignment is a minimum: never go below the ABI alignment of
+    // the allocated type, which is what an alloca without an explicit
+    // alignment gets.
+    unsigned alignment = 0;
+    if (std::optional<std::int64_t> requested = alloc.getAlignment())
+      alignment = static_cast<unsigned>(std::max<std::uint64_t>(
+          *requested, getDataLayout().getTypeABIAlignment(llvmObjectType)));
+
     // NOTE: we used to pass alloc->getAttrs() in the builder for non opaque
     // pointers! Only propagate pinned and bindc_name to help debugging, but
     // this should have no functional purpose (and passing the operand segment
     // attribute like before is certainly bad).
     auto llvmAlloc = mlir::LLVM::AllocaOp::create(
         rewriter, loc, ::getLlvmPtrType(alloc.getContext(), allocaAs),
-        llvmObjectType, size);
+        llvmObjectType, size, alignment);
     if (alloc.getPinned())
       llvmAlloc->setDiscardableAttr(alloc.getPinnedAttrName(),
                                     alloc.getPinnedAttr());
@@ -3848,8 +3856,11 @@ struct GlobalOpConversion : public fir::FIROpConversion<fir::GlobalOp> {
         0, getGlobalAddressSpace(rewriter), false,
         mlir::LLVM::ThreadLocalMode::NotThreadLocal, comdat, attrs, dbgExprs);
 
+    // A requested alignment is a minimum: never go below the ABI alignment of
+    // the global's type.
     if (global.getAlignment() && *global.getAlignment() > 0)
-      g.setAlignment(*global.getAlignment());
+      g.setAlignment(std::max<std::uint64_t>(
+          *global.getAlignment(), getDataLayout().getTypeABIAlignment(tyAttr)));
     if (auto section = global.getSection())
       g.setSection(*section);
 
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 3d079afa0d1a51..97bca969e96b7b 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -408,7 +408,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
                           llvm::ArrayRef<mlir::NamedAttribute> attributes) {
   auto nameAttr = builder.getStringAttr(uniqName);
   build(builder, result, wrapAllocaResultType(inType), inType, nameAttr, {},
-        /*pinned=*/false, typeparams, shape);
+        /*pinned=*/false, typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -419,7 +419,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
                           llvm::ArrayRef<mlir::NamedAttribute> attributes) {
   auto nameAttr = builder.getStringAttr(uniqName);
   build(builder, result, wrapAllocaResultType(inType), inType, nameAttr, {},
-        pinned, typeparams, shape);
+        pinned, typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -433,7 +433,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
   auto bindcAttr =
       bindcName.empty() ? mlir::StringAttr{} : builder.getStringAttr(bindcName);
   build(builder, result, wrapAllocaResultType(inType), inType, nameAttr,
-        bindcAttr, /*pinned=*/false, typeparams, shape);
+        bindcAttr, /*pinned=*/false, typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -448,7 +448,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
   auto bindcAttr =
       bindcName.empty() ? mlir::StringAttr{} : builder.getStringAttr(bindcName);
   build(builder, result, wrapAllocaResultType(inType), inType, nameAttr,
-        bindcAttr, pinned, typeparams, shape);
+        bindcAttr, pinned, typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -457,7 +457,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
                           mlir::ValueRange typeparams, mlir::ValueRange shape,
                           llvm::ArrayRef<mlir::NamedAttribute> attributes) {
   build(builder, result, wrapAllocaResultType(inType), inType, {}, {},
-        /*pinned=*/false, typeparams, shape);
+        /*pinned=*/false, typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -467,7 +467,7 @@ void fir::AllocaOp::build(mlir::OpBuilder &builder,
                           mlir::ValueRange shape,
                           llvm::ArrayRef<mlir::NamedAttribute> attributes) {
   build(builder, result, wrapAllocaResultType(inType), inType, {}, {}, pinned,
-        typeparams, shape);
+        typeparams, shape, /*alignment=*/{});
   result.addAttributes(attributes);
 }
 
@@ -492,6 +492,10 @@ llvm::LogicalResult fir::AllocaOp::verify() {
   mlir::Type outType = getType();
   if (!mlir::isa<fir::ReferenceType>(outType))
     return emitOpError("must be a !fir.ref type");
+  if (std::optional<std::int64_t> alignment = getAlignment()) {
+    if (*alignment <= 0 || (*alignment & (*alignment - 1)) != 0)
+      return emitOpError("alignment must be a positive power of two");
+  }
   return mlir::success();
 }
 
diff --git a/flang/test/Fir/convert-to-llvm.fir b/flang/test/Fir/convert-to-llvm.fir
index 73b930b3803285..a0c4d0f6db0f46 100644
--- a/flang/test/Fir/convert-to-llvm.fir
+++ b/flang/test/Fir/convert-to-llvm.fir
@@ -1115,6 +1115,38 @@ func.func @alloca_several() -> !fir.ref<i32> {
 
 // -----
 
+// Test fir.alloca with a requested alignment
+
+func.func @alloca_aligned() -> !fir.ref<!fir.array<10xf32>> {
+  %1 = fir.alloca !fir.array<10xf32> {alignment = 64 : i64}
+  return %1 : !fir.ref<!fir.array<10xf32>>
+}
+
+// CHECK-LABEL: llvm.func @alloca_aligned() -> !llvm.ptr
+// CHECK: [[N:%.*]] = llvm.mlir.constant(1 : i64) : i64
+// GENERIC: [[A:%.*]] = llvm.alloca [[N]] x !llvm.array<10 x f32> {alignment = 64 : i64}
+// AMDGPU: [[AA:%.*]] = llvm.alloca [[N]] x !llvm.array<10 x f32> {alignment = 64 : i64} : (i64) -> !llvm.ptr<5>
+// AMDGPU: [[A:%.*]] = llvm.addrspacecast [[AA]] : !llvm.ptr<5> to !llvm.ptr
+// CHECK: llvm.return [[A]] : !llvm.ptr
+
+// -----
+
+// A requested alignment below the ABI alignment of the type is raised to it.
+
+func.func @alloca_underaligned() -> !fir.ref<i32> {
+  %1 = fir.alloca i32 {alignment = 2 : i64}
+  return %1 : !fir.ref<i32>
+}
+
+// CHECK-LABEL: llvm.func @alloca_underaligned() -> !llvm.ptr
+// CHECK: [[N:%.*]] = llvm.mlir.constant(1 : i64) : i64
+// GENERIC: [[A:%.*]] = llvm.alloca [[N]] x i32 {alignment = 4 : i64}
+// AMDGPU: [[AA:%.*]] = llvm.alloca [[N]] x i32 {alignment = 4 : i64} : (i64) -> !llvm.ptr<5>
+// AMDGPU: [[A:%.*]] = llvm.addrspacecast [[AA]] : !llvm.ptr<5> to !llvm.ptr
+// CHECK: llvm.return [[A]] : !llvm.ptr
+
+// -----
+
 // Test fir.alloca of pointer to array
 
 func.func @alloca_ptr_to_array() -> !fir.ref<!fir.ptr<!fir.array<?xi32>>> {
diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir
index 704cd2ab217fa6..bf42a9c1ae79e1 100644
--- a/flang/test/Fir/fir-ops.fir
+++ b/flang/test/Fir/fir-ops.fir
@@ -37,10 +37,12 @@ func.func @instructions() {
 // CHECK: [[VAL_0:%.*]] = fir.alloca !fir.array<10xi32>
 // CHECK: [[VAL_1:%.*]] = fir.load [[VAL_0]] : !fir.ref<!fir.array<10xi32>>
 // CHECK: [[VAL_2:%.*]] = fir.alloca i32
+// CHECK: [[VAL_ALIGNED:%.*]] = fir.alloca !fir.array<10xf32> <{alignment = 64 : i64}>
 // CHECK: [[VAL_3:%.*]] = arith.constant 22 : i32
   %0 = fir.alloca !fir.array<10xi32>
   %1 = fir.load %0 : !fir.ref<!fir.array<10xi32>>
   %2 = fir.alloca i32
+  %aligned = fir.alloca !fir.array<10xf32> {alignment = 64 : i64}
   %3 = arith.constant 22 : i32
 
 // CHECK: fir.store [[VAL_3]] to [[VAL_2]] : !fir.ref<i32>
diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir
index 8964df1d2e93e6..8bf414f86173f9 100644
--- a/flang/test/Fir/invalid.fir
+++ b/flang/test/Fir/invalid.fir
@@ -1400,3 +1400,11 @@ func.func @allocmem_bad_alignment() -> !fir.heap<!fir.array<10xf32>> {
   %0 = fir.allocmem !fir.array<10xf32> {alignment = 3 : i64}
   return %0 : !fir.heap<!fir.array<10xf32>>
 }
+
+// -----
+
+func.func @alloca_bad_alignment() -> !fir.ref<!fir.array<10xf32>> {
+  // expected-error at +1 {{'fir.alloca' op alignment must be a positive power of two}}
+  %0 = fir.alloca !fir.array<10xf32> {alignment = 24 : i64}
+  return %0 : !fir.ref<!fir.array<10xf32>>
+}

>From ef50769176c1ba9d520ae422a551ef49e6e3ffb3 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 2/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}



More information about the flang-commits mailing list