[flang-commits] [flang] e3cd88a - [flang] Fixed StackArrays assertion after #121919. (#122550)
via flang-commits
flang-commits at lists.llvm.org
Mon Jan 13 11:56:15 PST 2025
Author: Slava Zakharin
Date: 2025-01-13T11:56:11-08:00
New Revision: e3cd88a7be1dfd912bb6e7c7e888e7b442ffb5de
URL: https://github.com/llvm/llvm-project/commit/e3cd88a7be1dfd912bb6e7c7e888e7b442ffb5de
DIFF: https://github.com/llvm/llvm-project/commit/e3cd88a7be1dfd912bb6e7c7e888e7b442ffb5de.diff
LOG: [flang] Fixed StackArrays assertion after #121919. (#122550)
`findAllocaLoopInsertionPoint()` hit assertion not being able
to find the `fir.freemem` because of the `fir.convert`.
I think it is better to look for `fir.freemem` same way
with the look-through walk.
Added:
Modified:
flang/lib/Optimizer/Transforms/StackArrays.cpp
flang/test/Transforms/stack-arrays.fir
Removed:
################################################################################
diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 2a9d3397e87b08..9a6566bef50f13 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -198,7 +198,9 @@ class AllocMemConversion : public mlir::OpRewritePattern<fir::AllocMemOp> {
/// 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);
+ static InsertionPoint
+ findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc,
+ const llvm::SmallVector<mlir::Operation *> &freeOps);
private:
/// Handle to the DFA (already run)
@@ -206,7 +208,9 @@ class AllocMemConversion : public mlir::OpRewritePattern<fir::AllocMemOp> {
/// 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);
+ 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>
@@ -484,6 +488,22 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
llvm::DenseSet<mlir::Value> freedValues;
point.appendFreedValues(freedValues);
+ // Find all fir.freemem operations corresponding to fir.allocmem
+ // in freedValues. It is best to find the association going back
+ // from fir.freemem to fir.allocmem through the def-use chains,
+ // so that we can use lookThroughDeclaresAndConverts same way
+ // the AllocationAnalysis is handling them.
+ llvm::DenseMap<mlir::Operation *, llvm::SmallVector<mlir::Operation *>>
+ allocToFreeMemMap;
+ func->walk([&](fir::FreeMemOp freeOp) {
+ mlir::Value memref = lookThroughDeclaresAndConverts(freeOp.getHeapref());
+ if (!freedValues.count(memref))
+ return;
+
+ auto allocMem = memref.getDefiningOp<fir::AllocMemOp>();
+ allocToFreeMemMap[allocMem].push_back(freeOp);
+ });
+
// We only replace allocations which are definately freed on all routes
// through the function because otherwise the allocation may have an intende
// lifetime longer than the current stack frame (e.g. a heap allocation which
@@ -491,7 +511,8 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
for (mlir::Value freedValue : freedValues) {
fir::AllocMemOp allocmem = freedValue.getDefiningOp<fir::AllocMemOp>();
InsertionPoint insertionPoint =
- AllocMemConversion::findAllocaInsertionPoint(allocmem);
+ AllocMemConversion::findAllocaInsertionPoint(
+ allocmem, allocToFreeMemMap[allocmem]);
if (insertionPoint)
candidateOps.insert({allocmem, insertionPoint});
}
@@ -578,8 +599,9 @@ static bool isInLoop(mlir::Operation *op) {
op->getParentOfType<mlir::LoopLikeOpInterface>();
}
-InsertionPoint
-AllocMemConversion::findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc) {
+InsertionPoint AllocMemConversion::findAllocaInsertionPoint(
+ fir::AllocMemOp &oldAlloc,
+ const llvm::SmallVector<mlir::Operation *> &freeOps) {
// Ideally the alloca should be inserted at the end of the function entry
// block so that we do not allocate stack space in a loop. However,
// the operands to the alloca may not be available that early, so insert it
@@ -596,7 +618,7 @@ AllocMemConversion::findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc) {
if (isInLoop(oldAllocOp)) {
// where we want to put it is in a loop, and even the old location is in
// a loop. Give up.
- return findAllocaLoopInsertionPoint(oldAlloc);
+ return findAllocaLoopInsertionPoint(oldAlloc, freeOps);
}
return {oldAllocOp};
}
@@ -657,28 +679,14 @@ AllocMemConversion::findAllocaInsertionPoint(fir::AllocMemOp &oldAlloc) {
return checkReturn(&entryBlock);
}
-InsertionPoint
-AllocMemConversion::findAllocaLoopInsertionPoint(fir::AllocMemOp &oldAlloc) {
+InsertionPoint AllocMemConversion::findAllocaLoopInsertionPoint(
+ fir::AllocMemOp &oldAlloc,
+ const llvm::SmallVector<mlir::Operation *> &freeOps) {
mlir::Operation *oldAllocOp = oldAlloc;
// This is only called as a last resort. We should try to insert at the
// location of the old allocation, which is inside of a loop, using
// llvm.stacksave/llvm.stackrestore
- // find freemem ops
- llvm::SmallVector<mlir::Operation *, 1> freeOps;
-
- for (mlir::Operation *user : oldAllocOp->getUsers()) {
- if (auto declareOp = mlir::dyn_cast_if_present<fir::DeclareOp>(user)) {
- for (mlir::Operation *user : declareOp->getUsers()) {
- if (mlir::isa<fir::FreeMemOp>(user))
- freeOps.push_back(user);
- }
- }
-
- if (mlir::isa<fir::FreeMemOp>(user))
- freeOps.push_back(user);
- }
-
assert(freeOps.size() && "DFA should only return freed memory");
// Don't attempt to reason about a stacksave/stackrestore between
diff erent
diff --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index 444136d53e0340..a784cea9bc3a46 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -418,3 +418,45 @@ func.func @lookthrough() {
// CHECK: func.func @lookthrough() {
// CHECK: fir.alloca !fir.array<42xi32>
// CHECK-NOT: fir.freemem
+
+// StackArrays is better to find fir.freemem ops corresponding to fir.allocmem
+// using the same look through mechanism as during the allocation analysis,
+// looking through fir.convert and fir.declare.
+func.func @finding_freemem_in_block() {
+ %c0 = arith.constant 0 : index
+ %c10_i32 = arith.constant 10 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32 {bindc_name = "k", uniq_name = "k"}
+ %1 = fir.declare %0 {uniq_name = "k"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ fir.store %c1_i32 to %1 : !fir.ref<i32>
+ cf.br ^bb1
+^bb1: // 2 preds: ^bb0, ^bb2
+ %2 = fir.load %1 : !fir.ref<i32>
+ %3 = arith.cmpi sle, %2, %c10_i32 : i32
+ cf.cond_br %3, ^bb2, ^bb3
+^bb2: // pred: ^bb1
+ %4 = fir.declare %1 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "x"} : (!fir.ref<i32>) -> !fir.ref<i32>
+ %5 = fir.load %4 : !fir.ref<i32>
+ %6 = fir.convert %5 : (i32) -> index
+ %7 = arith.cmpi sgt, %6, %c0 : index
+ %8 = arith.select %7, %6, %c0 : index
+ %9 = fir.shape %8 : (index) -> !fir.shape<1>
+ %10 = fir.allocmem !fir.array<?xi32>, %8 {bindc_name = ".tmp.expr_result", uniq_name = ""}
+ %11 = fir.convert %10 : (!fir.heap<!fir.array<?xi32>>) -> !fir.ref<!fir.array<?xi32>>
+ %12 = fir.declare %11(%9) {uniq_name = ".tmp.expr_result"} : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xi32>>
+ %13 = fir.embox %12(%9) : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<?xi32>>
+ %14 = fir.call @_QPfunc(%1) fastmath<fast> : (!fir.ref<i32>) -> !fir.array<?xi32>
+ fir.save_result %14 to %12(%9) : !fir.array<?xi32>, !fir.ref<!fir.array<?xi32>>, !fir.shape<1>
+ fir.call @_QPsub(%13) fastmath<fast> : (!fir.box<!fir.array<?xi32>>) -> ()
+ %15 = fir.convert %12 : (!fir.ref<!fir.array<?xi32>>) -> !fir.heap<!fir.array<?xi32>>
+ fir.freemem %15 : !fir.heap<!fir.array<?xi32>>
+ %16 = fir.load %1 : !fir.ref<i32>
+ %17 = arith.addi %16, %c1_i32 : i32
+ fir.store %17 to %1 : !fir.ref<i32>
+ cf.br ^bb1
+^bb3: // pred: ^bb1
+ return
+}
+// CHECK: func.func @finding_freemem_in_block() {
+// CHECK: fir.alloca !fir.array<?xi32>
+// CHECK-NOT: fir.freemem
More information about the flang-commits
mailing list