[flang-commits] [flang] [flang][cuf] Support memory finalization at a final return statement (PR #116304)

[via flang-commits]

https://github.com/khaki3 updated https://github.com/llvm/llvm-project/pull/116304

>From 798f1bc5ad4e0b8e2e591ac5eadd6b09b8d4700f Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 14 Nov 2024 16:17:57 -0800
Subject: [PATCH 1/2] [flang][cuf] Support memory finalization at early exit
 points

---
 flang/include/flang/Lower/StatementContext.h | 11 +++---
 flang/lib/Lower/Bridge.cpp                   | 37 +++++++++-----------
 flang/test/Lower/CUDA/cuda-return.cuf        | 14 ++++++++
 3 files changed, 38 insertions(+), 24 deletions(-)
 create mode 100644 flang/test/Lower/CUDA/cuda-return.cuf

diff --git a/flang/include/flang/Lower/StatementContext.h b/flang/include/flang/Lower/StatementContext.h
index 7776edc93ed737..4a07eeaefece41 100644
--- a/flang/include/flang/Lower/StatementContext.h
+++ b/flang/include/flang/Lower/StatementContext.h
@@ -79,23 +79,26 @@ class StatementContext {
     }
   }
 
-  /// Make cleanup calls. Retain the stack top list for a repeat call.
+  /// Make a cleanup call. Retain the stack top list for a repeat call.
   void finalizeAndKeep() {
     assert(!cufs.empty() && "invalid finalize statement context");
     if (cufs.back())
       (*cufs.back())();
   }
 
-  /// Make cleanup calls. Clear the stack top list.
+  /// Make a cleanup call. Clear the stack top list.
   void finalizeAndReset() {
     finalizeAndKeep();
     cufs.back().reset();
   }
 
-  /// Make cleanup calls. Pop the stack top list.
+  /// Pop the stack top list.
+  void pop() { cufs.pop_back(); }
+
+  /// Make a cleanup call. Pop the stack top list.
   void finalizeAndPop() {
     finalizeAndKeep();
-    cufs.pop_back();
+    pop();
   }
 
   bool hasCode() const {
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index da53edf7e734b0..7f41742bf5e8b2 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -1621,13 +1621,19 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   // Termination of symbolically referenced execution units
   //===--------------------------------------------------------------------===//
 
-  /// END of program
+  /// Exit of a routine
   ///
-  /// Generate the cleanup block before the program exits
-  void genExitRoutine() {
-
-    if (blockIsUnterminated())
-      builder->create<mlir::func::ReturnOp>(toLocation());
+  /// Generate the cleanup block before the routine exits
+  void genExitRoutine(bool earlyReturn, mlir::ValueRange retval = {}) {
+    if (blockIsUnterminated()) {
+      bridge.openAccCtx().finalizeAndKeep();
+      bridge.fctCtx().finalizeAndKeep();
+      builder->create<mlir::func::ReturnOp>(toLocation(), retval);
+    }
+    if (!earlyReturn) {
+      bridge.openAccCtx().pop();
+      bridge.fctCtx().pop();
+    }
   }
 
   /// END of procedure-like constructs
@@ -1684,9 +1690,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
             resultRef = builder->createConvert(loc, resultRefType, resultRef);
           return builder->create<fir::LoadOp>(loc, resultRef);
         });
-    bridge.openAccCtx().finalizeAndPop();
-    bridge.fctCtx().finalizeAndPop();
-    builder->create<mlir::func::ReturnOp>(loc, resultVal);
+    genExitRoutine(false, resultVal);
   }
 
   /// Get the return value of a call to \p symbol, which is a subroutine entry
@@ -1712,13 +1716,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     } else if (Fortran::semantics::HasAlternateReturns(symbol)) {
       mlir::Value retval = builder->create<fir::LoadOp>(
           toLocation(), getAltReturnResult(symbol));
-      bridge.openAccCtx().finalizeAndPop();
-      bridge.fctCtx().finalizeAndPop();
-      builder->create<mlir::func::ReturnOp>(toLocation(), retval);
+      genExitRoutine(false, retval);
     } else {
-      bridge.openAccCtx().finalizeAndPop();
-      bridge.fctCtx().finalizeAndPop();
-      genExitRoutine();
+      genExitRoutine(false);
     }
   }
 
@@ -5018,8 +5018,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
       it->stmtCtx.finalizeAndKeep();
     }
     if (funit->isMainProgram()) {
-      bridge.fctCtx().finalizeAndKeep();
-      genExitRoutine();
+      genExitRoutine(true);
       return;
     }
     mlir::Location loc = toLocation();
@@ -5478,9 +5477,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   void endNewFunction(Fortran::lower::pft::FunctionLikeUnit &funit) {
     setCurrentPosition(Fortran::lower::pft::stmtSourceLoc(funit.endStmt));
     if (funit.isMainProgram()) {
-      bridge.openAccCtx().finalizeAndPop();
-      bridge.fctCtx().finalizeAndPop();
-      genExitRoutine();
+      genExitRoutine(false);
     } else {
       genFIRProcedureExit(funit, funit.getSubprogramSymbol());
     }
diff --git a/flang/test/Lower/CUDA/cuda-return.cuf b/flang/test/Lower/CUDA/cuda-return.cuf
new file mode 100644
index 00000000000000..b40d63aa8b5fe0
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-return.cuf
@@ -0,0 +1,14 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! Check if the the finalizer works
+
+program main
+  integer, device :: a(10)
+  return
+end
+
+! CHECK: func.func @_QQmain() attributes {fir.bindc_name = "main"} {
+! CHECK: %[[DECL:.*]]:2 = hlfir.declare
+! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
+! CHECK-NEXT: return
+! CHECK-NEXT: }

>From 172149c1468d5ff22c4030b6bd286fb0101a6c41 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Thu, 14 Nov 2024 17:11:32 -0800
Subject: [PATCH 2/2] [test] Fix a comment

---
 flang/test/Lower/CUDA/cuda-return.cuf | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flang/test/Lower/CUDA/cuda-return.cuf b/flang/test/Lower/CUDA/cuda-return.cuf
index b40d63aa8b5fe0..c9f9a8b57ef041 100644
--- a/flang/test/Lower/CUDA/cuda-return.cuf
+++ b/flang/test/Lower/CUDA/cuda-return.cuf
@@ -1,6 +1,6 @@
 ! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
 
-! Check if the the finalizer works
+! Check if finalization works with a return statement
 
 program main
   integer, device :: a(10)