[flang-commits] [flang] [flang][cuda] Do not trigger automatic deallocation in main (PR #128789)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/128789

Similar to host flow, do not trigger automatic deallocation at then end of the main program since anything could happen like a cudaDevcieReset(). 

>From 4bf0ac2fdf1ef75f20745e0b7902c5cdd0aef45b Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 25 Feb 2025 15:45:41 -0800
Subject: [PATCH] [flang][cuda] Do not trigger automatic deallocation in main

F
---
 flang/lib/Lower/ConvertVariable.cpp     | 15 ++++++++-----
 flang/test/Lower/CUDA/cuda-return01.cuf | 29 ++++++++++++++++++++-----
 flang/test/Lower/CUDA/cuda-return02.cuf |  2 --
 3 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp
index 81d14fbb1d777..cc55191170c65 100644
--- a/flang/lib/Lower/ConvertVariable.cpp
+++ b/flang/lib/Lower/ConvertVariable.cpp
@@ -977,12 +977,15 @@ static void instantiateLocal(Fortran::lower::AbstractConverter &converter,
     fir::ExtendedValue exv =
         converter.getSymbolExtendedValue(var.getSymbol(), &symMap);
     auto *sym = &var.getSymbol();
-    converter.getFctCtx().attachCleanup([builder, loc, exv, sym]() {
-      cuf::DataAttributeAttr dataAttr =
-          Fortran::lower::translateSymbolCUFDataAttribute(builder->getContext(),
-                                                          *sym);
-      builder->create<cuf::FreeOp>(loc, fir::getBase(exv), dataAttr);
-    });
+    const Fortran::semantics::Scope &owner = sym->owner();
+    if (owner.kind() != Fortran::semantics::Scope::Kind::MainProgram) {
+      converter.getFctCtx().attachCleanup([builder, loc, exv, sym]() {
+        cuf::DataAttributeAttr dataAttr =
+            Fortran::lower::translateSymbolCUFDataAttribute(
+                builder->getContext(), *sym);
+        builder->create<cuf::FreeOp>(loc, fir::getBase(exv), dataAttr);
+      });
+    }
   }
   if (std::optional<VariableCleanUp> cleanup =
           needDeallocationOrFinalization(var)) {
diff --git a/flang/test/Lower/CUDA/cuda-return01.cuf b/flang/test/Lower/CUDA/cuda-return01.cuf
index c9f9a8b57ef04..0fb1dbe63fe2d 100644
--- a/flang/test/Lower/CUDA/cuda-return01.cuf
+++ b/flang/test/Lower/CUDA/cuda-return01.cuf
@@ -2,13 +2,32 @@
 
 ! Check if finalization works with a return statement
 
-program main
+subroutine sub1
   integer, device :: a(10)
   return
 end
 
-! CHECK: func.func @_QQmain() attributes {fir.bindc_name = "main"} {
+! CHECK: func.func @_QPsub1()
 ! CHECK: %[[DECL:.*]]:2 = hlfir.declare
-! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
-! CHECK-NEXT: return
-! CHECK-NEXT: }
+! CHECK: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
+! CHECK: return
+! CHECK: }
+
+subroutine sub2
+  integer, device, allocatable :: a(:)
+  return
+end
+
+! CHECK-LABEL: func.func @_QPsub2()
+! CHECK: fir.if
+! CHECK: cuf.deallocate 
+! CHECK: cuf.free
+
+program main
+  integer, allocatable, device :: a(:)
+  return
+end
+
+! CHECK-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "main"}
+! CHECK: cuf.alloc !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+! CHECK-NOT: cuf.free
diff --git a/flang/test/Lower/CUDA/cuda-return02.cuf b/flang/test/Lower/CUDA/cuda-return02.cuf
index 5d01f0a24b420..9a0fa14f8973d 100644
--- a/flang/test/Lower/CUDA/cuda-return02.cuf
+++ b/flang/test/Lower/CUDA/cuda-return02.cuf
@@ -17,10 +17,8 @@ end
 ! CHECK: %[[DECL:.*]]:2 = hlfir.declare
 ! CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 ! CHECK-NEXT: ^bb1:
-! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
 ! CHECK-NEXT: return
 ! CHECK-NEXT: ^bb2:
-! CHECK-NEXT: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<10xi32>>
 ! CHECK-NEXT: return
 ! CHECK-NEXT: }