[flang-commits] [flang] [flang] add some missing stackrestore with -fstack-arrays (PR #208161)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/208161

In the StackArray pass that moves array temporaries from the heap to the stack under -fstack-arrays, when visiting FreeMemOp to insert stackrestore, the code was not unwrapping converts as done in other parts of the code leading to the allocation conversion and stacksave insertion to happen without the emission of the stackrestore.

Reuse the same utility as in the rest of the pass to get consistent behavior and fix the memory leak.

>From b3e16a8248c4050a2fd1ee1784dd21290a1c3635 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 8 Jul 2026 00:24:12 -0700
Subject: [PATCH] [flang] add some missing stackrestore with -fstack-arrays

---
 .../lib/Optimizer/Transforms/StackArrays.cpp  | 16 ++++------
 flang/test/Transforms/stack-arrays.fir        | 30 +++++++++++++++++++
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index cfbd273b88ea6..4af26415a39d2 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -804,17 +804,11 @@ AllocMemConversion::insertAlloca(fir::AllocMemOp &oldAlloc,
 static void
 visitFreeMemOp(fir::AllocMemOp oldAlloc,
                const std::function<void(mlir::Operation *)> &callBack) {
-  for (mlir::Operation *user : oldAlloc->getUsers()) {
-    if (auto declareOp = mlir::dyn_cast_if_present<fir::DeclareOp>(user)) {
-      for (mlir::Operation *user : declareOp->getUsers()) {
-        if (mlir::isa<fir::FreeMemOp>(user))
-          callBack(user);
-      }
-    }
-
-    if (mlir::isa<fir::FreeMemOp>(user))
-      callBack(user);
-  }
+  mlir::Operation *parent = oldAlloc->getParentOp();
+  parent->walk([&](fir::FreeMemOp freeOp) {
+    if (lookThroughDeclaresAndConverts(freeOp->getOperand(0)) == oldAlloc)
+      callBack(freeOp);
+  });
 }
 
 void AllocMemConversion::insertStackSaveRestore(
diff --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index 25fc73153003a..7a005c69a1b04 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -495,5 +495,35 @@ func.func @finding_freemem_in_block() {
   return
 }
 // CHECK-LABEL: func.func @finding_freemem_in_block()
+// CHECK:     llvm.intr.stacksave
 // CHECK:     fir.alloca !fir.array<?xi32>
+// CHECK:     llvm.intr.stackrestore
 // CHECK-NOT: fir.freemem
+
+// Test temporary inside a CFG loop where freemem is reached via a convert
+// chain.
+func.func private @_QPuse_array(!fir.ref<!fir.array<?xf64>>) -> ()
+func.func @freemem_through_convert_in_loop(%arg0: i1, %n: index) {
+  cf.br ^bb1
+^bb1:  // 2 preds: ^bb0, ^bb2
+  cf.cond_br %arg0, ^bb2, ^bb3
+^bb2:  // pred: ^bb1
+  %shape = fir.shape %n : (index) -> !fir.shape<1>
+  %mem = fir.allocmem !fir.array<?xf64>, %n {bindc_name = ".tmp.array", uniq_name = ""}
+  %decl = fir.declare %mem(%shape) {uniq_name = ".tmp.array"} : (!fir.heap<!fir.array<?xf64>>, !fir.shape<1>) -> !fir.heap<!fir.array<?xf64>>
+  %ref = fir.convert %decl : (!fir.heap<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+  fir.call @_QPuse_array(%ref) : (!fir.ref<!fir.array<?xf64>>) -> ()
+  %heap = fir.convert %ref : (!fir.ref<!fir.array<?xf64>>) -> !fir.heap<!fir.array<?xf64>>
+  fir.freemem %heap : !fir.heap<!fir.array<?xf64>>
+  cf.br ^bb1
+^bb3:  // pred: ^bb1
+  return
+}
+// CHECK-LABEL: func.func @freemem_through_convert_in_loop(
+// CHECK:       ^bb2:
+// CHECK:         %[[SP:.*]] = llvm.intr.stacksave : !llvm.ptr
+// CHECK:         fir.alloca !fir.array<?xf64>
+// CHECK:         fir.call @_QPuse_array(
+// CHECK:         llvm.intr.stackrestore %[[SP]] : !llvm.ptr
+// CHECK-NOT:     fir.freemem
+// CHECK:         cf.br ^bb1