[llvm-branch-commits] [flang] [flang] Add policy-driven allocation-placement pass - memory passes unification [2/5] (PR #210742)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Jul 27 07:32:07 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
Author: jeanPerier
<details>
<summary>Changes</summary>
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
---
Patch is 27.27 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/210742.diff
8 Files Affected:
- (added) flang/include/flang/Optimizer/Transforms/AllocationPlacementPolicy.h (+90)
- (modified) flang/include/flang/Optimizer/Transforms/Passes.h (+8)
- (modified) flang/include/flang/Optimizer/Transforms/Passes.td (+26)
- (added) flang/lib/Optimizer/Transforms/AllocationPlacement.cpp (+324)
- (modified) flang/lib/Optimizer/Transforms/CMakeLists.txt (+1)
- (added) flang/test/Transforms/allocation-placement-budget.fir (+29)
- (added) flang/test/Transforms/allocation-placement-stack-arrays-mode.fir (+24)
- (added) flang/test/Transforms/allocation-placement.fir (+84)
``````````diff
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 different 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. different 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 8573c0e4f3f00..8ef0d33a12a82 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()) {...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/210742
More information about the llvm-branch-commits
mailing list