[flang-commits] [flang] [flang][cuda] Set alloca block in cuf kernel (PR #128776)

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

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/128776

>From 1b355fec5325e749403b267ae20427b805823d65 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Tue, 25 Feb 2025 13:50:00 -0800
Subject: [PATCH 1/2] [flang][cuda] Set alloca block in cuf kernel

---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp    |  3 ++
 .../Lower/CUDA/cuda-kernel-alloca-block.cuf   | 29 +++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 flang/test/Lower/CUDA/cuda-kernel-alloca-block.cuf

diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index d9779c46ae79e..d6eb6d9336881 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -264,6 +264,9 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
     return recipeIface.getAllocaBlock(getRegion());
   }
 
+  if (auto cufKernelOp = getRegion().getParentOfType<cuf::KernelOp>())
+    return &cufKernelOp.getRegion().front();
+
   return getEntryBlock();
 }
 
diff --git a/flang/test/Lower/CUDA/cuda-kernel-alloca-block.cuf b/flang/test/Lower/CUDA/cuda-kernel-alloca-block.cuf
new file mode 100644
index 0000000000000..afad65bb3138c
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-kernel-alloca-block.cuf
@@ -0,0 +1,29 @@
+! RUN: bbc -emit-fir -hlfir -fcuda %s -o - | FileCheck %s
+
+module char1
+
+character*1, dimension(:), allocatable, device :: da, db
+
+contains
+  attributes(device) logical function check_char( c1, c2 )
+    character*1, value :: c1, c2
+    check_char = .false.
+    if(c1 .eq. c2) check_char = .true.
+  end function
+end module
+
+program p
+  use char1
+  integer :: n, i
+  !$cuf kernel do(1)<<<*,*>>>
+  do i = 1, n
+    if (check_char(da(i),db(i))) then
+      print*,'same'
+    endif
+  enddo
+end program
+
+! CHECK-LABEL: func.func @_QQmain()
+! CHECK: cuf.kernel
+! CHECK-COUNT-2: %{{.*}} = fir.alloca !fir.char<1> {bindc_name = ".tmp"}
+

>From b3408cbd9c31ebe728294041f8c0b3a9c0682012 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 2/2] [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: }