[flang-commits] [flang] 005c18e - [flang] Add policy-driven allocation-placement pass - memory passes unification [2/5] (#210742)
via flang-commits
flang-commits at lists.llvm.org
Wed Jul 29 02:01:54 PDT 2026
Author: jeanPerier
Date: 2026-07-29T11:01:47+02:00
New Revision: 005c18eac0a70158c2db4611bf824af6e0bc80ac
URL: https://github.com/llvm/llvm-project/commit/005c18eac0a70158c2db4611bf824af6e0bc80ac
DIFF: https://github.com/llvm/llvm-project/commit/005c18eac0a70158c2db4611bf824af6e0bc80ac.diff
LOG: [flang] Add policy-driven allocation-placement pass - memory passes unification [2/5] (#210742)
Introduce a new function-level pass, allocation-placement, that unifies
the stack/heap placement decisions currently split between the
stack-arrays and memory-allocation-opt passes. For each array allocation
it consults a policy to decide whether it should live on the stack
(fir.alloca) or the heap (fir.allocmem) and rewrites it accordingly,
reusing fir::replaceAllocas for stack-to-heap and the StackArrays
analysis/rewrite for heap-to-stack (so heap-to-stack only happens where
it is provably safe).
The default policy (AllocationPlacementPolicy.h) is threshold-driven:
- small constant-size arrays go on the stack within a per-function stack
budget, otherwise on the heap;
- big constant-size arrays: user variables stay on the stack,
temporaries go on the heap;
- runtime-sized arrays go on the heap;
- an aggressive mode places all arrays on the stack (best effort). User
variables are distinguished from compiler temporaries via the presence
of a uniqued name. A hook lets downstream users override the thresholds
per allocation (e.g. for device routines or parallel regions).
The pass is not wired into any pipeline yet; it is reachable through
fir-opt and covered by isolated tests.
Assisted-by: AI
Added:
flang/include/flang/Optimizer/Transforms/AllocationPlacementPolicy.h
flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
flang/test/Transforms/allocation-placement-budget.fir
flang/test/Transforms/allocation-placement-stack-arrays-mode.fir
flang/test/Transforms/allocation-placement.fir
Modified:
flang/include/flang/Optimizer/Transforms/Passes.h
flang/include/flang/Optimizer/Transforms/Passes.td
flang/lib/Optimizer/Transforms/CMakeLists.txt
Removed:
################################################################################
diff --git a/flang/include/flang/Optimizer/Transforms/AllocationPlacementPolicy.h b/flang/include/flang/Optimizer/Transforms/AllocationPlacementPolicy.h
new file mode 100644
index 0000000000000..001e4cdfe06cc
--- /dev/null
+++ b/flang/include/flang/Optimizer/Transforms/AllocationPlacementPolicy.h
@@ -0,0 +1,90 @@
+//===-- Optimizer/Transforms/AllocationPlacementPolicy.h --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header defines the policy used by the allocation-placement pass to
+// decide whether an array allocation should live on the stack (fir.alloca) or
+// on the heap (fir.allocmem). The policy is expressed with tunable thresholds
+// and a per-function stack budget, and can be customized through a hook (e.g.
+// with
diff erent thresholds inside device routines or parallel regions).
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H
+#define FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <optional>
+
+namespace mlir {
+class Operation;
+} // namespace mlir
+
+namespace fir {
+
+/// Desired placement for an array allocation.
+enum class AllocationPlacement {
+ /// The allocation should live on the stack (fir.alloca).
+ Stack,
+ /// The allocation should live on the heap (fir.allocmem).
+ Heap,
+ /// The allocation should be left where it currently is.
+ Leave,
+};
+
+/// Tunable thresholds controlling where array allocations are placed.
+struct AllocationPlacementThresholds {
+ /// Place all array allocations on the stack when possible (-fstack-arrays).
+ /// When false, use the size/kind-based placement policy.
+ bool stackArrays = false;
+ /// Constant-size arrays up to this many bytes are considered "small".
+ std::size_t smallArrayThresholdBytes = 64;
+ /// Per-function budget (in bytes) for small arrays placed on the stack.
+ std::size_t totalStackLimitBytes = 4ull * 1024 * 1024;
+};
+
+/// Facts about a single array allocation used to decide its placement.
+struct AllocationInfo {
+ /// The allocation operation (fir.alloca or fir.allocmem).
+ mlir::Operation *op = nullptr;
+ /// True if the allocation currently lives on the stack (fir.alloca).
+ bool isCurrentlyOnStack = false;
+ /// True if the allocation is a compiler temporary (as opposed to a user
+ /// variable).
+ bool isTemporary = false;
+ /// True if the allocation has a runtime-determined size. Note this is not the
+ /// same as !byteSize: a constant-size array may have no computable byteSize
+ /// (e.g. when no data layout is available), in which case it is not dynamic
+ /// but its size is still unknown.
+ bool isDynamic = false;
+ /// The constant size of the allocation in bytes, if it can be determined.
+ std::optional<std::int64_t> byteSize;
+};
+
+/// Default placement policy. Decides where \p info should live given \p
+/// thresholds and the per-function stack bytes already committed to the stack
+/// (\p stackBytesUsed). The caller is responsible for updating \p
+/// stackBytesUsed based on the returned decision.
+AllocationPlacement
+decideAllocationPlacement(const AllocationInfo &info,
+ const AllocationPlacementThresholds &thresholds,
+ std::size_t stackBytesUsed);
+
+/// Placement decision hook. Has the same signature as decideAllocationPlacement
+/// so the policy can be fully overridden (e.g.
diff erent thresholds inside
+/// device routines or parallel regions); a hook may adjust the thresholds and
+/// delegate to decideAllocationPlacement.
+using AllocationPlacementHook = std::function<AllocationPlacement(
+ const AllocationInfo & /*info*/,
+ const AllocationPlacementThresholds & /*thresholds*/,
+ std::size_t /*stackBytesUsed*/)>;
+
+} // namespace fir
+
+#endif // FORTRAN_OPTIMIZER_TRANSFORMS_ALLOCATIONPLACEMENTPOLICY_H
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
index 0dc6182187aa3..8b77fb9918f4e 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.h
+++ b/flang/include/flang/Optimizer/Transforms/Passes.h
@@ -9,6 +9,7 @@
#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_PASSES_H
#define FORTRAN_OPTIMIZER_TRANSFORMS_PASSES_H
+#include "flang/Optimizer/Transforms/AllocationPlacementPolicy.h"
#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
@@ -41,6 +42,13 @@ enum class LICMNestedHoistingMode {
#include "flang/Optimizer/Transforms/Passes.h.inc"
+/// Create the allocation-placement pass with the given options and a hook that
+/// can override the thresholds per allocation (e.g. for device routines or
+/// parallel regions). This complements the tablegen-generated overloads.
+std::unique_ptr<mlir::Pass>
+createAllocationPlacement(const AllocationPlacementOptions &options,
+ AllocationPlacementHook placementHook);
+
std::unique_ptr<mlir::Pass> createAffineDemotionPass();
std::unique_ptr<mlir::Pass>
createArrayValueCopyPass(fir::ArrayValueCopyOptions options = {});
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index e51ac93596546..ba29cd6df2eac 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -312,6 +312,32 @@ def MemoryAllocationOpt : Pass<"memory-allocation-opt", "mlir::func::FuncOp"> {
];
}
+def AllocationPlacement : Pass<"allocation-placement", "mlir::func::FuncOp"> {
+ let summary = "Place array allocations on the stack or the heap by policy.";
+ let description = [{
+ Unified array allocation placement. Converts fir.alloca to fir.allocmem and
+ fir.allocmem to fir.alloca based on tunable byte-size thresholds, a
+ per-function stack budget, and whether the allocation is a user variable or
+ a compiler temporary. Heap-to-stack conversions are only performed where the
+ allocation is provably freed on all paths through the function.
+ }];
+ let dependentDialects = [
+ "fir::FIROpsDialect", "mlir::DLTIDialect", "mlir::LLVM::LLVMDialect",
+ "mlir::arith::ArithDialect"
+ ];
+ let options = [
+ Option<"stackArrays", "stack-arrays", "bool", /*default=*/"false",
+ "Place all array allocations on the stack (-fstack-arrays); "
+ "otherwise use the size/kind-based placement policy.">,
+ Option<"smallArrayThresholdBytes", "small-array-threshold", "std::size_t",
+ /*default=*/"64",
+ "Constant-size arrays up to this many bytes are considered small.">,
+ Option<"totalStackLimitBytes", "total-stack-limit", "std::size_t",
+ /*default=*/"4194304",
+ "Per-function budget (bytes) for small arrays placed on the stack.">
+ ];
+}
+
// This needs to be a "mlir::ModuleOp" pass, because it inserts global constants
def ConstantArgumentGlobalisationOpt : Pass<"constant-argument-globalisation-opt", "mlir::ModuleOp"> {
let summary = "Convert constant function arguments to global constants.";
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
new file mode 100644
index 0000000000000..162cb5fbb21fe
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -0,0 +1,324 @@
+//===- AllocationPlacement.cpp --------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass decides, for each array allocation in a function, whether it should
+// live on the stack (fir.alloca) or on the heap (fir.allocmem), and rewrites it
+// accordingly. The decision is delegated to the policy in
+// AllocationPlacementPolicy.h. Two rewrite engines are reused:
+// - stack-to-heap uses fir::replaceAllocas (MemoryUtils);
+// - heap-to-stack reuses the StackArrays analysis and rewrite, which only
+// stackifies fir.allocmem that are provably freed on all paths.
+//
+//===----------------------------------------------------------------------===//
+
+#include "StackArrays.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "flang/Optimizer/Support/DataLayout.h"
+#include "flang/Optimizer/Transforms/AllocationPlacementPolicy.h"
+#include "flang/Optimizer/Transforms/MemoryUtils.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+#include <optional>
+
+namespace fir {
+#define GEN_PASS_DEF_ALLOCATIONPLACEMENT
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "allocation-placement"
+
+//===----------------------------------------------------------------------===//
+// Default placement policy
+//===----------------------------------------------------------------------===//
+
+fir::AllocationPlacement
+fir::decideAllocationPlacement(const AllocationInfo &info,
+ const AllocationPlacementThresholds &thresholds,
+ std::size_t stackBytesUsed) {
+ using P = fir::AllocationPlacement;
+
+ // Translate a "should this be on the stack" decision into a placement,
+ // accounting for where the allocation currently lives.
+ auto place = [&](bool wantStack) -> P {
+ if (wantStack)
+ return info.isCurrentlyOnStack ? P::Leave : P::Stack;
+ return info.isCurrentlyOnStack ? P::Heap : P::Leave;
+ };
+
+ // -fstack-arrays: put everything on the stack (best effort). The
+ // heap-to-stack conversion still only happens where it is provably safe.
+ if (thresholds.stackArrays)
+ return place(/*wantStack=*/true);
+
+ // Runtime-sized arrays (automatic arrays, dynamic temporaries) go on the
+ // heap.
+ if (info.isDynamic)
+ return place(/*wantStack=*/false);
+
+ // Without a known constant size we cannot reason about thresholds.
+ if (!info.byteSize)
+ return P::Leave;
+
+ // Constant-size user variables always go on the stack.
+ if (!info.isTemporary)
+ return place(/*wantStack=*/true);
+
+ auto size = static_cast<std::size_t>(*info.byteSize);
+ if (size <= thresholds.smallArrayThresholdBytes)
+ // Small arrays go on the stack while the per-function budget allows it.
+ return place(/*wantStack=*/stackBytesUsed + size <=
+ thresholds.totalStackLimitBytes);
+
+ // Big array temporaries go on the heap.
+ return place(/*wantStack=*/false);
+}
+
+namespace {
+
+/// Return true if the allocation is a compiler temporary, i.e. it has no
+/// uniqued name (user variables always carry one).
+static bool isTemporaryAllocation(mlir::Operation *op) {
+ std::optional<llvm::StringRef> uniqName;
+ if (auto alloca = mlir::dyn_cast<fir::AllocaOp>(op))
+ uniqName = alloca.getUniqName();
+ else if (auto allocmem = mlir::dyn_cast<fir::AllocMemOp>(op))
+ uniqName = allocmem.getUniqName();
+ return !uniqName || uniqName->empty();
+}
+
+/// Return the constant byte size of an array allocation, or std::nullopt if it
+/// is dynamic or cannot be determined.
+static std::optional<std::int64_t>
+getConstantByteSize(mlir::Operation *op,
+ const std::optional<mlir::DataLayout> &dl,
+ const std::optional<fir::KindMapping> &kindMap) {
+ if (!dl || !kindMap)
+ return std::nullopt;
+ if (auto alloca = mlir::dyn_cast<fir::AllocaOp>(op))
+ return fir::getAllocaByteSize(alloca, *dl, *kindMap);
+ if (auto allocmem = mlir::dyn_cast<fir::AllocMemOp>(op)) {
+ if (allocmem.hasLenParams() || allocmem.hasShapeOperands())
+ return std::nullopt;
+ if (auto sizeAndAlignment = fir::getTypeSizeAndAlignment(
+ op->getLoc(), allocmem.getAllocatedType(), *dl, *kindMap))
+ return static_cast<std::int64_t>(sizeAndAlignment->first);
+ }
+ return std::nullopt;
+}
+
+/// Replacement generator used for stack-to-heap conversions (fir.alloca ->
+/// fir.allocmem). Mirrors the MemoryAllocation pass.
+static mlir::Value genAllocmem(mlir::OpBuilder &builder, fir::AllocaOp alloca,
+ bool /*deallocPointsDominateAlloc*/) {
+ mlir::Type varTy = alloca.getInType();
+ auto unpackName = [](std::optional<llvm::StringRef> opt) -> llvm::StringRef {
+ if (opt)
+ return *opt;
+ return {};
+ };
+ llvm::StringRef uniqName = unpackName(alloca.getUniqName());
+ llvm::StringRef bindcName = unpackName(alloca.getBindcName());
+ auto heap = fir::AllocMemOp::create(builder, alloca.getLoc(), varTy, uniqName,
+ bindcName, alloca.getTypeparams(),
+ alloca.getShape());
+ LLVM_DEBUG(llvm::dbgs() << "allocation placement: replaced " << alloca
+ << " with " << heap << '\n');
+ return heap;
+}
+
+static void genFreemem(mlir::Location loc, mlir::OpBuilder &builder,
+ mlir::Value allocmem) {
+ [[maybe_unused]] auto free = fir::FreeMemOp::create(builder, loc, allocmem);
+ LLVM_DEBUG(llvm::dbgs() << "allocation placement: add free " << free
+ << " for " << allocmem << '\n');
+}
+
+[[maybe_unused]] static llvm::StringRef
+placementName(fir::AllocationPlacement placement) {
+ switch (placement) {
+ case fir::AllocationPlacement::Stack:
+ return "stack";
+ case fir::AllocationPlacement::Heap:
+ return "heap";
+ case fir::AllocationPlacement::Leave:
+ return "leave";
+ }
+ return "?";
+}
+
+/// Return true if \p placement leaves the allocation on the stack.
+static bool endsUpOnStack(fir::AllocationPlacement placement,
+ bool isCurrentlyOnStack) {
+ return placement == fir::AllocationPlacement::Stack ||
+ (placement == fir::AllocationPlacement::Leave && isCurrentlyOnStack);
+}
+
+class AllocationPlacementPass
+ : public fir::impl::AllocationPlacementBase<AllocationPlacementPass> {
+public:
+ AllocationPlacementPass() = default;
+ AllocationPlacementPass(const AllocationPlacementPass &pass)
+ : fir::impl::AllocationPlacementBase<AllocationPlacementPass>(pass),
+ placementHook(pass.placementHook) {}
+ AllocationPlacementPass(fir::AllocationPlacementOptions options)
+ : fir::impl::AllocationPlacementBase<AllocationPlacementPass>(
+ std::move(options)) {}
+ AllocationPlacementPass(fir::AllocationPlacementOptions options,
+ fir::AllocationPlacementHook hook)
+ : fir::impl::AllocationPlacementBase<AllocationPlacementPass>(
+ std::move(options)),
+ placementHook(std::move(hook)) {}
+
+ void runOnOperation() override;
+
+private:
+ fir::AllocationPlacementHook placementHook;
+};
+
+void AllocationPlacementPass::runOnOperation() {
+ mlir::func::FuncOp func = getOperation();
+ if (func.empty())
+ return;
+
+ fir::AllocationPlacementThresholds baseThresholds;
+ baseThresholds.stackArrays = stackArrays;
+ baseThresholds.smallArrayThresholdBytes = smallArrayThresholdBytes;
+ baseThresholds.totalStackLimitBytes = totalStackLimitBytes;
+
+ auto module = func->getParentOfType<mlir::ModuleOp>();
+ std::optional<mlir::DataLayout> dl =
+ module ? fir::support::getOrSetMLIRDataLayout(
+ module, /*allowDefaultLayout=*/false)
+ : std::nullopt;
+ std::optional<fir::KindMapping> kindMap;
+ if (module)
+ kindMap = fir::getKindMapping(module);
+
+ // Analysis of which fir.allocmem can be safely moved to the stack, and where.
+ auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
+ const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
+ analysis.getCandidateOps(func);
+ if (!candidateOps) {
+ signalPassFailure();
+ return;
+ }
+
+ // Walk allocations in deterministic program order, maintaining the running
+ // per-function stack budget while collecting the conversions to perform.
+ std::size_t stackBytesUsed = 0;
+ llvm::DenseSet<mlir::Operation *> allocasToHeap;
+ llvm::SmallVector<mlir::Operation *> allocmemsToStack;
+
+ func.walk([&](mlir::Operation *op) {
+ auto alloca = mlir::dyn_cast<fir::AllocaOp>(op);
+ auto allocmem = mlir::dyn_cast<fir::AllocMemOp>(op);
+ if (!alloca && !allocmem)
+ return;
+
+ // Only array allocations are considered.
+ mlir::Type inTy = alloca ? alloca.getInType() : allocmem.getAllocatedType();
+ if (!mlir::isa<fir::SequenceType>(inTy))
+ return;
+
+ fir::AllocationInfo info;
+ info.op = op;
+ info.isCurrentlyOnStack = static_cast<bool>(alloca);
+ info.isTemporary = isTemporaryAllocation(op);
+ info.isDynamic =
+ alloca ? alloca.isDynamic()
+ : (allocmem.hasLenParams() || allocmem.hasShapeOperands());
+ info.byteSize = getConstantByteSize(op, dl, kindMap);
+
+ // A hook, if provided, fully overrides the default policy; it may delegate
+ // back to decideAllocationPlacement after adjusting the thresholds.
+ fir::AllocationPlacement placement =
+ placementHook ? placementHook(info, baseThresholds, stackBytesUsed)
+ : fir::decideAllocationPlacement(info, baseThresholds,
+ stackBytesUsed);
+
+ // Account for the decision in the running stack budget.
+ if (endsUpOnStack(placement, info.isCurrentlyOnStack) && info.byteSize)
+ stackBytesUsed += static_cast<std::size_t>(*info.byteSize);
+
+ LLVM_DEBUG({
+ llvm::dbgs() << "allocation-placement: "
+ << (info.isCurrentlyOnStack ? "alloca" : "allocmem") << " "
+ << (info.isTemporary ? "temp" : "user") << " ";
+ if (info.isDynamic)
+ llvm::dbgs() << "dynamic";
+ else if (info.byteSize)
+ llvm::dbgs() << *info.byteSize << "B";
+ else
+ llvm::dbgs() << "unknown-size";
+ llvm::dbgs() << " -> " << placementName(placement) << "\n";
+ });
+
+ if (alloca && placement == fir::AllocationPlacement::Heap)
+ allocasToHeap.insert(op);
+ else if (allocmem && placement == fir::AllocationPlacement::Stack &&
+ candidateOps->contains(op))
+ allocmemsToStack.push_back(op);
+ });
+
+ LLVM_DEBUG(llvm::dbgs() << "allocation-placement: " << func.getSymName()
+ << ": " << allocasToHeap.size() << " stack->heap, "
+ << allocmemsToStack.size() << " heap->stack, "
+ << stackBytesUsed << " stack bytes used\n");
+
+ // Heap-to-stack: only provably-safe candidates are converted.
+ if (!allocmemsToStack.empty()) {
+ mlir::MLIRContext &context = getContext();
+ mlir::RewritePatternSet patterns(&context);
+ mlir::GreedyRewriteConfig config;
+ config.setRegionSimplificationLevel(
+ mlir::GreedySimplifyRegionLevel::Disabled);
+ config.setStrictness(mlir::GreedyRewriteStrictness::ExistingAndNewOps);
+ config.enableFolding(false);
+ patterns.insert<fir::AllocMemConversion>(&context, *candidateOps, dl,
+ kindMap);
+ if (mlir::failed(mlir::applyOpPatternsGreedily(
+ allocmemsToStack, std::move(patterns), config))) {
+ mlir::emitError(func->getLoc(),
+ "error in allocation placement (heap to stack)\n");
+ signalPassFailure();
+ return;
+ }
+ }
+
+ // Stack-to-heap.
+ if (!allocasToHeap.empty()) {
+ mlir::IRRewriter rewriter(&getContext());
+ auto mustReplace = [&](fir::AllocaOp alloca) {
+ return allocasToHeap.contains(alloca.getOperation());
+ };
+ fir::replaceAllocas(rewriter, func.getOperation(), mustReplace, genAllocmem,
+ genFreemem);
+ }
+}
+
+} // namespace
+
+std::unique_ptr<mlir::Pass>
+fir::createAllocationPlacement(const fir::AllocationPlacementOptions &options,
+ fir::AllocationPlacementHook placementHook) {
+ return std::make_unique<AllocationPlacementPass>(options,
+ std::move(placementHook));
+}
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 1a4940af95d3e..997dc22063138 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -5,6 +5,7 @@ add_flang_library(FIRTransforms
AffineDemotion.cpp
AffinePromotion.cpp
AlgebraicSimplification.cpp
+ AllocationPlacement.cpp
AnnotateConstant.cpp
ArrayValueCopy.cpp
ArrayValueCopy.cpp
diff --git a/flang/test/Transforms/allocation-placement-budget.fir b/flang/test/Transforms/allocation-placement-budget.fir
new file mode 100644
index 0000000000000..6da58c4cfe9be
--- /dev/null
+++ b/flang/test/Transforms/allocation-placement-budget.fir
@@ -0,0 +1,29 @@
+// Test the per-function stack budget: small arrays are placed on the stack
+// until the total stack limit is reached, after which further small arrays are
+// left on the heap. With a 64-byte limit, the first 40-byte array fits on the
+// stack but the following 48-byte array does not (40 + 48 = 88 > 64).
+
+// RUN: fir-opt --allocation-placement="small-array-threshold=64 total-stack-limit=64" %s | FileCheck %s
+
+module attributes {fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"} {
+
+// CHECK-LABEL: func.func @budget
+// CHECK-DAG: fir.alloca !fir.array<10xi32>
+// CHECK-DAG: fir.allocmem !fir.array<12xi32>
+func.func @budget() {
+ %0 = fir.allocmem !fir.array<10xi32>
+ %1 = fir.allocmem !fir.array<12xi32>
+ %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>
+ %r1 = fir.convert %1 : (!fir.heap<!fir.array<12xi32>>) -> !fir.ref<!fir.array<12xi32>>
+ %e1 = fir.coordinate_of %r1, %c0 : (!fir.ref<!fir.array<12xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e1 : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<10xi32>>
+ fir.freemem %1 : !fir.heap<!fir.array<12xi32>>
+ return
+}
+
+}
diff --git a/flang/test/Transforms/allocation-placement-stack-arrays-mode.fir b/flang/test/Transforms/allocation-placement-stack-arrays-mode.fir
new file mode 100644
index 0000000000000..fb0d7afaea62a
--- /dev/null
+++ b/flang/test/Transforms/allocation-placement-stack-arrays-mode.fir
@@ -0,0 +1,24 @@
+// Test the aggressive (-fstack-arrays) allocation-placement mode: all array
+// temporaries are moved to the stack regardless of size, as long as the
+// heap-to-stack conversion is provably safe.
+
+// RUN: fir-opt --allocation-placement="stack-arrays=true" %s | FileCheck %s
+
+module attributes {fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"} {
+
+// Big temporary heap allocation -> stack under -fstack-arrays.
+// CHECK-LABEL: func.func @big_temp_to_stack
+// CHECK: fir.alloca !fir.array<100xi32>
+// CHECK-NOT: fir.allocmem
+func.func @big_temp_to_stack() {
+ %0 = fir.allocmem !fir.array<100xi32>
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 0 : i32
+ %r = fir.convert %0 : (!fir.heap<!fir.array<100xi32>>) -> !fir.ref<!fir.array<100xi32>>
+ %e = fir.coordinate_of %r, %c0 : (!fir.ref<!fir.array<100xi32>>, index) -> !fir.ref<i32>
+ fir.store %v to %e : !fir.ref<i32>
+ fir.freemem %0 : !fir.heap<!fir.array<100xi32>>
+ return
+}
+
+}
diff --git a/flang/test/Transforms/allocation-placement.fir b/flang/test/Transforms/allocation-placement.fir
new file mode 100644
index 0000000000000..ef8fce82c03ac
--- /dev/null
+++ b/flang/test/Transforms/allocation-placement.fir
@@ -0,0 +1,84 @@
+// Test the default (-fno-stack-arrays) allocation-placement policy:
+// - small constant-size arrays go on the stack (within the budget),
+// - big constant-size arrays: user variables stay on the stack, temporaries
+// go on the heap,
+// - runtime-sized arrays go on the heap.
+// A user variable is identified by a non-empty uniq_name; a temporary has none.
+// i32 is 4 bytes, so <10xi32> = 40 bytes (small) and <100xi32> = 400 bytes (big)
+// with the default 64-byte small threshold.
+
+// RUN: fir-opt --allocation-placement %s | FileCheck %s
+
+module attributes {fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"} {
+
+// Small temporary heap allocation -> stack.
+// CHECK-LABEL: func.func @small_temp
+// CHECK: fir.alloca !fir.array<10xi32>
+// CHECK-NOT: fir.allocmem
+func.func @small_temp() {
+ %0 = fir.allocmem !fir.array<10xi32>
+ %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
+}
+
+// Small temporary stack allocation -> stays on the stack.
+// CHECK-LABEL: func.func @small_temp_alloca
+// CHECK: fir.alloca !fir.array<10xi32>
+// CHECK-NOT: fir.allocmem
+func.func @small_temp_alloca() {
+ %0 = fir.alloca !fir.array<10xi32>
+ return
+}
+
+// Big temporary heap allocation -> stays on the heap.
+// CHECK-LABEL: func.func @big_temp
+// CHECK: fir.allocmem !fir.array<100xi32>
+// CHECK: fir.freemem
+func.func @big_temp() {
+ %0 = fir.allocmem !fir.array<100xi32>
+ fir.freemem %0 : !fir.heap<!fir.array<100xi32>>
+ return
+}
+
+// Big user-variable stack allocation -> stays on the stack.
+// CHECK-LABEL: func.func @big_user
+// CHECK: fir.alloca !fir.array<100xi32>
+// CHECK-NOT: fir.allocmem
+func.func @big_user() {
+ %0 = fir.alloca !fir.array<100xi32> {bindc_name = "arr", uniq_name = "_QFbig_userEarr"}
+ return
+}
+
+// Big temporary stack allocation -> heap.
+// CHECK-LABEL: func.func @big_temp_alloca
+// CHECK: fir.allocmem !fir.array<100xi32>
+// CHECK: fir.freemem
+func.func @big_temp_alloca() {
+ %0 = fir.alloca !fir.array<100xi32>
+ return
+}
+
+// Runtime-sized user variable (automatic array) -> heap.
+// CHECK-LABEL: func.func @dyn_user
+// CHECK: fir.allocmem !fir.array<?xi32>
+// CHECK: fir.freemem
+func.func @dyn_user(%n: index) {
+ %0 = fir.alloca !fir.array<?xi32>, %n {bindc_name = "arr", uniq_name = "_QFdyn_userEarr"}
+ return
+}
+
+// Runtime-sized temporary -> stays on the heap.
+// CHECK-LABEL: func.func @dyn_temp
+// CHECK: fir.allocmem !fir.array<?xi32>
+func.func @dyn_temp(%n: index) {
+ %0 = fir.allocmem !fir.array<?xi32>, %n
+ fir.freemem %0 : !fir.heap<!fir.array<?xi32>>
+ return
+}
+
+}
More information about the flang-commits
mailing list