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

[via flang-commits]

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

We generate `cuf.free` and `func.return` twice if a return statement exists at the end of program.

```f90
program test
  integer, device :: a(10)
  return
end
```

```
% flang -x cuda test.cuf -mmlir --mlir-print-ir-after-all
error: loc("/path/to/test.cuf":3:3): 'func.return' op must be the last operation in the parent block
// -----// IR Dump After Fortran::lower::VerifierPass Failed () //----- //
```

Dumped IR:
```mlir
  "func.func"() <{function_type = () -> (), sym_name = "_QQmain"}> ({
...
    "cuf.free"(%5#1) <{data_attr = #cuf.cuda<device>}> : (!fir.ref<!fir.array<10xi32>>) -> ()
    "func.return"() : () -> ()
    "cuf.free"(%5#1) <{data_attr = #cuf.cuda<device>}> : (!fir.ref<!fir.array<10xi32>>) -> ()
    "func.return"() : () -> ()
}
...
```

We have `genExitRoutine` in `Bridge.cpp` to generate `func.return`, which is guarded by `blockIsUnterminated()`. However, we redundantly runs `bridge.fctCtx().finalizeAndKeep()` in this case, resulting in two pairs of `cuf.free` and `func.return`. This PR fixes `Bridge.cpp` by using `blockIsUnterminated()` to guard `finalizeAndKeep` as well.

>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] [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: }