[flang][cuda] Set alloca block in cuf kernel

[clementval]

Temporary created during lowering in a cuf kernel must be set in the cuf kernel itself otherwise they will be allocated on the host.

[llvmbot]

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

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

Changes

Temporary created during lowering in a cuf kernel must be set in the cuf kernel itself otherwise they will be allocated on the host.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/Builder/FIRBuilder.cpp (+3)
• (added) flang/test/Lower/CUDA/cuda-kernel-alloca-block.cuf (+29)

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"}
+