[flang-commits] [flang] [flang][NFC] Extract StackArrays analysis and rewrite into a header - memory passes unification [1/6] (PR #210721)
via flang-commits
flang-commits at lists.llvm.org
Mon Jul 20 06:48:53 PDT 2026
https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/210721
Move InsertionPoint, StackArraysAnalysisWrapper, and AllocMemConversion out of the anonymous namespace in StackArrays.cpp into a new StackArrays.h header (in namespace fir), so the "which fir.allocmem can be safely moved to the stack, and where" analysis and the heap-to-stack rewrite pattern can be reused by other passes.
The dataflow internals (AllocationState, LatticePoint, AllocationAnalysis), the command-line options, and all method definitions remain in the .cpp. No functional change intended.
>From e9a17d03184d4389987a55e3eb6a306b71c74918 Mon Sep 17 00:00:00 2001
From: root <jperier at nvidia.com>
Date: Fri, 17 Jul 2026 09:00:11 -0700
Subject: [PATCH] [flang][NFC] Extract StackArrays analysis and rewrite into a
header
Move InsertionPoint, StackArraysAnalysisWrapper, and AllocMemConversion
out of the anonymous namespace in StackArrays.cpp into a new
StackArrays.h header (in namespace fir), so the "which fir.allocmem can
be safely moved to the stack, and where" analysis and the heap-to-stack
rewrite pattern can be reused by other passes.
The dataflow internals (AllocationState, LatticePoint, AllocationAnalysis),
the command-line options, and all method definitions remain in the .cpp.
No functional change intended.
---
.../lib/Optimizer/Transforms/StackArrays.cpp | 145 ++---------------
flang/lib/Optimizer/Transforms/StackArrays.h | 154 ++++++++++++++++++
2 files changed, 170 insertions(+), 129 deletions(-)
create mode 100644 flang/lib/Optimizer/Transforms/StackArrays.h
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 4af26415a39d2..77861e67a07b1 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include "StackArrays.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/LowLevelIntrinsics.h"
#include "flang/Optimizer/Dialect/FIRAttr.h"
@@ -74,52 +75,6 @@ enum class AllocationState {
Allocated,
};
-/// Stores where an alloca should be inserted. If the PointerUnion is an
-/// Operation the alloca should be inserted /after/ the operation. If it is a
-/// block, the alloca can be placed anywhere in that block.
-class InsertionPoint {
- llvm::PointerUnion<mlir::Operation *, mlir::Block *> location;
- bool saveRestoreStack;
-
- /// Get contained pointer type or nullptr
- template <class T>
- T *tryGetPtr() const {
- // Use llvm::dyn_cast_if_present because location may be null here.
- if (T *ptr = llvm::dyn_cast_if_present<T *>(location))
- return ptr;
- return nullptr;
- }
-
-public:
- template <class T>
- InsertionPoint(T *ptr, bool saveRestoreStack = false)
- : location(ptr), saveRestoreStack{saveRestoreStack} {}
- InsertionPoint(std::nullptr_t null)
- : location(null), saveRestoreStack{false} {}
-
- /// Get contained operation, or nullptr
- mlir::Operation *tryGetOperation() const {
- return tryGetPtr<mlir::Operation>();
- }
-
- /// Get contained block, or nullptr
- mlir::Block *tryGetBlock() const { return tryGetPtr<mlir::Block>(); }
-
- /// Get whether the stack should be saved/restored. If yes, an llvm.stacksave
- /// intrinsic should be added before the alloca, and an llvm.stackrestore
- /// intrinsic should be added where the freemem is
- bool shouldSaveRestoreStack() const { return saveRestoreStack; }
-
- operator bool() const { return tryGetOperation() || tryGetBlock(); }
-
- bool operator==(const InsertionPoint &rhs) const {
- return (location == rhs.location) &&
- (saveRestoreStack == rhs.saveRestoreStack);
- }
-
- bool operator!=(const InsertionPoint &rhs) const { return !(*this == rhs); }
-};
-
/// Maps SSA values to their AllocationState at a particular program point.
/// Also caches the insertion points for the new alloca operations
class LatticePoint : public mlir::dataflow::AbstractDenseLattice {
@@ -174,74 +129,6 @@ class AllocationAnalysis
mlir::LogicalResult processOperation(mlir::Operation *op) override;
};
-/// Drives analysis to find candidate fir.allocmem operations which could be
-/// moved to the stack. Intended to be used with mlir::Pass::getAnalysis
-class StackArraysAnalysisWrapper {
-public:
- MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StackArraysAnalysisWrapper)
-
- // Maps fir.allocmem -> place to insert alloca
- using AllocMemMap = llvm::DenseMap<mlir::Operation *, InsertionPoint>;
-
- StackArraysAnalysisWrapper(mlir::Operation *op) {}
-
- // returns nullptr if analysis failed
- const AllocMemMap *getCandidateOps(mlir::Operation *func);
-
-private:
- llvm::DenseMap<mlir::Operation *, AllocMemMap> funcMaps;
-
- llvm::LogicalResult analyseFunction(mlir::Operation *func);
-};
-
-/// Converts a fir.allocmem to a fir.alloca
-class AllocMemConversion : public mlir::OpRewritePattern<fir::AllocMemOp> {
-public:
- explicit AllocMemConversion(
- mlir::MLIRContext *ctx,
- const StackArraysAnalysisWrapper::AllocMemMap &candidateOps,
- std::optional<mlir::DataLayout> &dl,
- std::optional<fir::KindMapping> &kindMap)
- : OpRewritePattern(ctx), candidateOps{candidateOps}, dl{dl},
- kindMap{kindMap} {}
-
- llvm::LogicalResult
- matchAndRewrite(fir::AllocMemOp allocmem,
- mlir::PatternRewriter &rewriter) const override;
-
- /// Determine where to insert the alloca operation. The returned value should
- /// be checked to see if it is inside a loop
- static InsertionPoint
- findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc,
- const llvm::SmallVector<mlir::Operation *> &freeOps);
-
-private:
- /// Handle to the DFA (already run)
- const StackArraysAnalysisWrapper::AllocMemMap &candidateOps;
-
- const std::optional<mlir::DataLayout> &dl;
- const std::optional<fir::KindMapping> &kindMap;
-
- /// If we failed to find an insertion point not inside a loop, see if it would
- /// be safe to use an llvm.stacksave/llvm.stackrestore inside the loop
- static InsertionPoint findAllocaLoopInsertionPoint(
- fir::AllocMemOp &oldAlloc,
- const llvm::SmallVector<mlir::Operation *> &freeOps);
-
- /// Returns the alloca if it was successfully inserted, otherwise {}
- std::optional<fir::AllocaOp>
- insertAlloca(fir::AllocMemOp &oldAlloc,
- mlir::PatternRewriter &rewriter) const;
-
- /// Inserts a stacksave before oldAlloc and a stackrestore after each freemem
- void insertStackSaveRestore(fir::AllocMemOp oldAlloc,
- mlir::PatternRewriter &rewriter) const;
- /// Emit lifetime markers for newAlloc between oldAlloc and each freemem.
- /// If the allocation is dynamic, no life markers are emitted.
- void insertLifetimeMarkers(fir::AllocMemOp oldAlloc, fir::AllocaOp newAlloc,
- mlir::PatternRewriter &rewriter) const;
-};
-
class StackArraysPass : public fir::impl::StackArraysBase<StackArraysPass> {
public:
StackArraysPass() = default;
@@ -462,7 +349,7 @@ mlir::LogicalResult AllocationAnalysis::processOperation(mlir::Operation *op) {
}
llvm::LogicalResult
-StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
+fir::StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
assert(mlir::isa<mlir::func::FuncOp>(func));
size_t nAllocs = 0;
func->walk([&nAllocs](fir::AllocMemOp) { nAllocs++; });
@@ -543,8 +430,8 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
return mlir::success();
}
-const StackArraysAnalysisWrapper::AllocMemMap *
-StackArraysAnalysisWrapper::getCandidateOps(mlir::Operation *func) {
+const fir::StackArraysAnalysisWrapper::AllocMemMap *
+fir::StackArraysAnalysisWrapper::getCandidateOps(mlir::Operation *func) {
if (!funcMaps.contains(func))
if (mlir::failed(analyseFunction(func)))
return nullptr;
@@ -575,9 +462,8 @@ static mlir::Value convertAllocationType(mlir::PatternRewriter &rewriter,
return conv;
}
-llvm::LogicalResult
-AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
- mlir::PatternRewriter &rewriter) const {
+llvm::LogicalResult fir::AllocMemConversion::matchAndRewrite(
+ fir::AllocMemOp allocmem, mlir::PatternRewriter &rewriter) const {
auto oldInsertionPt = rewriter.saveInsertionPoint();
// add alloca operation
std::optional<fir::AllocaOp> alloca = insertAlloca(allocmem, rewriter);
@@ -615,7 +501,7 @@ static bool isInLoop(mlir::Operation *op) {
op->getParentOfType<mlir::LoopLikeOpInterface>();
}
-InsertionPoint AllocMemConversion::findAllocaInsertionPoint(
+fir::InsertionPoint fir::AllocMemConversion::findAllocaInsertionPoint(
fir::AllocMemOp &oldAlloc,
const llvm::SmallVector<mlir::Operation *> &freeOps) {
// Ideally the alloca should be inserted at the end of the function entry
@@ -731,7 +617,7 @@ InsertionPoint AllocMemConversion::findAllocaInsertionPoint(
return checkReturn(&entryBlock);
}
-InsertionPoint AllocMemConversion::findAllocaLoopInsertionPoint(
+fir::InsertionPoint fir::AllocMemConversion::findAllocaLoopInsertionPoint(
fir::AllocMemOp &oldAlloc,
const llvm::SmallVector<mlir::Operation *> &freeOps) {
mlir::Operation *oldAllocOp = oldAlloc;
@@ -762,8 +648,8 @@ InsertionPoint AllocMemConversion::findAllocaLoopInsertionPoint(
}
std::optional<fir::AllocaOp>
-AllocMemConversion::insertAlloca(fir::AllocMemOp &oldAlloc,
- mlir::PatternRewriter &rewriter) const {
+fir::AllocMemConversion::insertAlloca(fir::AllocMemOp &oldAlloc,
+ mlir::PatternRewriter &rewriter) const {
auto it = candidateOps.find(oldAlloc.getOperation());
if (it == candidateOps.end())
return {};
@@ -811,7 +697,7 @@ visitFreeMemOp(fir::AllocMemOp oldAlloc,
});
}
-void AllocMemConversion::insertStackSaveRestore(
+void fir::AllocMemConversion::insertStackSaveRestore(
fir::AllocMemOp oldAlloc, mlir::PatternRewriter &rewriter) const {
mlir::OpBuilder::InsertionGuard insertGuard(rewriter);
auto mod = oldAlloc->getParentOfType<mlir::ModuleOp>();
@@ -827,7 +713,7 @@ void AllocMemConversion::insertStackSaveRestore(
visitFreeMemOp(oldAlloc, createStackRestoreCall);
}
-void AllocMemConversion::insertLifetimeMarkers(
+void fir::AllocMemConversion::insertLifetimeMarkers(
fir::AllocMemOp oldAlloc, fir::AllocaOp newAlloc,
mlir::PatternRewriter &rewriter) const {
if (!dl || !kindMap)
@@ -860,8 +746,8 @@ llvm::StringRef StackArraysPass::getDescription() const {
void StackArraysPass::runOnOperation() {
mlir::func::FuncOp func = getOperation();
- auto &analysis = getAnalysis<StackArraysAnalysisWrapper>();
- const StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
+ auto &analysis = getAnalysis<fir::StackArraysAnalysisWrapper>();
+ const fir::StackArraysAnalysisWrapper::AllocMemMap *candidateOps =
analysis.getCandidateOps(func);
if (!candidateOps) {
signalPassFailure();
@@ -893,7 +779,8 @@ void StackArraysPass::runOnOperation() {
if (module)
kindMap = fir::getKindMapping(module);
- patterns.insert<AllocMemConversion>(&context, *candidateOps, dl, kindMap);
+ patterns.insert<fir::AllocMemConversion>(&context, *candidateOps, dl,
+ kindMap);
if (mlir::failed(mlir::applyOpPatternsGreedily(
opsToConvert, std::move(patterns), config))) {
mlir::emitError(func->getLoc(), "error in stack arrays optimization\n");
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.h b/flang/lib/Optimizer/Transforms/StackArrays.h
new file mode 100644
index 0000000000000..f6a0805a7de78
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/StackArrays.h
@@ -0,0 +1,154 @@
+//===- StackArrays.h ------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+//
+// This header exposes the reusable pieces of the StackArrays pass: the
+// analysis that determines which fir.allocmem operations can safely be moved
+// to the stack (and where the replacement fir.alloca should be inserted), and
+// the pattern that performs the heap-to-stack rewrite. They are shared so that
+// other passes can reuse the same "is this heap allocation safely stackifiable,
+// and where" logic.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_STACKARRAYS_H
+#define FORTRAN_OPTIMIZER_TRANSFORMS_STACKARRAYS_H
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Support/TypeID.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include <optional>
+
+namespace fir {
+
+/// Stores where an alloca should be inserted. If the PointerUnion is an
+/// Operation the alloca should be inserted /after/ the operation. If it is a
+/// block, the alloca can be placed anywhere in that block.
+class InsertionPoint {
+ llvm::PointerUnion<mlir::Operation *, mlir::Block *> location;
+ bool saveRestoreStack;
+
+ /// Get contained pointer type or nullptr
+ template <class T>
+ T *tryGetPtr() const {
+ // Use llvm::dyn_cast_if_present because location may be null here.
+ if (T *ptr = llvm::dyn_cast_if_present<T *>(location))
+ return ptr;
+ return nullptr;
+ }
+
+public:
+ template <class T>
+ InsertionPoint(T *ptr, bool saveRestoreStack = false)
+ : location(ptr), saveRestoreStack{saveRestoreStack} {}
+ InsertionPoint(std::nullptr_t null)
+ : location(null), saveRestoreStack{false} {}
+
+ /// Get contained operation, or nullptr
+ mlir::Operation *tryGetOperation() const {
+ return tryGetPtr<mlir::Operation>();
+ }
+
+ /// Get contained block, or nullptr
+ mlir::Block *tryGetBlock() const { return tryGetPtr<mlir::Block>(); }
+
+ /// Get whether the stack should be saved/restored. If yes, an llvm.stacksave
+ /// intrinsic should be added before the alloca, and an llvm.stackrestore
+ /// intrinsic should be added where the freemem is
+ bool shouldSaveRestoreStack() const { return saveRestoreStack; }
+
+ operator bool() const { return tryGetOperation() || tryGetBlock(); }
+
+ bool operator==(const InsertionPoint &rhs) const {
+ return (location == rhs.location) &&
+ (saveRestoreStack == rhs.saveRestoreStack);
+ }
+
+ bool operator!=(const InsertionPoint &rhs) const { return !(*this == rhs); }
+};
+
+/// Drives analysis to find candidate fir.allocmem operations which could be
+/// moved to the stack. Intended to be used with mlir::Pass::getAnalysis
+class StackArraysAnalysisWrapper {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(StackArraysAnalysisWrapper)
+
+ // Maps fir.allocmem -> place to insert alloca
+ using AllocMemMap = llvm::DenseMap<mlir::Operation *, InsertionPoint>;
+
+ StackArraysAnalysisWrapper(mlir::Operation *op) {}
+
+ // returns nullptr if analysis failed
+ const AllocMemMap *getCandidateOps(mlir::Operation *func);
+
+private:
+ llvm::DenseMap<mlir::Operation *, AllocMemMap> funcMaps;
+
+ llvm::LogicalResult analyseFunction(mlir::Operation *func);
+};
+
+/// Converts a fir.allocmem to a fir.alloca
+class AllocMemConversion : public mlir::OpRewritePattern<fir::AllocMemOp> {
+public:
+ explicit AllocMemConversion(
+ mlir::MLIRContext *ctx,
+ const StackArraysAnalysisWrapper::AllocMemMap &candidateOps,
+ std::optional<mlir::DataLayout> &dl,
+ std::optional<fir::KindMapping> &kindMap)
+ : OpRewritePattern(ctx), candidateOps{candidateOps}, dl{dl},
+ kindMap{kindMap} {}
+
+ llvm::LogicalResult
+ matchAndRewrite(fir::AllocMemOp allocmem,
+ mlir::PatternRewriter &rewriter) const override;
+
+ /// Determine where to insert the alloca operation. The returned value should
+ /// be checked to see if it is inside a loop
+ static InsertionPoint
+ findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc,
+ const llvm::SmallVector<mlir::Operation *> &freeOps);
+
+private:
+ /// Handle to the DFA (already run)
+ const StackArraysAnalysisWrapper::AllocMemMap &candidateOps;
+
+ const std::optional<mlir::DataLayout> &dl;
+ const std::optional<fir::KindMapping> &kindMap;
+
+ /// If we failed to find an insertion point not inside a loop, see if it would
+ /// be safe to use an llvm.stacksave/llvm.stackrestore inside the loop
+ static InsertionPoint findAllocaLoopInsertionPoint(
+ fir::AllocMemOp &oldAlloc,
+ const llvm::SmallVector<mlir::Operation *> &freeOps);
+
+ /// Returns the alloca if it was successfully inserted, otherwise {}
+ std::optional<fir::AllocaOp>
+ insertAlloca(fir::AllocMemOp &oldAlloc,
+ mlir::PatternRewriter &rewriter) const;
+
+ /// Inserts a stacksave before oldAlloc and a stackrestore after each freemem
+ void insertStackSaveRestore(fir::AllocMemOp oldAlloc,
+ mlir::PatternRewriter &rewriter) const;
+ /// Emit lifetime markers for newAlloc between oldAlloc and each freemem.
+ /// If the allocation is dynamic, no life markers are emitted.
+ void insertLifetimeMarkers(fir::AllocMemOp oldAlloc, fir::AllocaOp newAlloc,
+ mlir::PatternRewriter &rewriter) const;
+};
+
+} // namespace fir
+
+#endif // FORTRAN_OPTIMIZER_TRANSFORMS_STACKARRAYS_H
More information about the flang-commits
mailing list