[flang][cuda] Do not trigger automatic deallocation in main

[clementval]

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

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/128789.diff

3 Files Affected:

• (modified) flang/lib/Lower/ConvertVariable.cpp (+9-6)
• (modified) flang/test/Lower/CUDA/cuda-return01.cuf (+24-5)
• (modified) flang/test/Lower/CUDA/cuda-return02.cuf (-2)

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