[flang][cuda] Do not check assignment semantic in cuf kernel

[clementval]

CUDA Fortran data transfer are not generated in cuf kernel. Avoid triggering false semantic error when in a cuf kernel.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

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

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

Changes

CUDA Fortran data transfer are not generated in cuf kernel. Avoid triggering false semantic error when in a cuf kernel.

---

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

3 Files Affected:

• (modified) flang/lib/Semantics/check-cuda.cpp (+6-1)
• (modified) flang/lib/Semantics/check-cuda.h (+2)
• (modified) flang/test/Lower/CUDA/cuda-data-transfer.cuf (+21)

diff --git a/flang/lib/Semantics/check-cuda.cpp b/flang/lib/Semantics/check-cuda.cpp
index 60b8b3213fdb82..d8bd435bf09f38 100644
--- a/flang/lib/Semantics/check-cuda.cpp
+++ b/flang/lib/Semantics/check-cuda.cpp
@@ -565,13 +565,18 @@ void CUDAChecker::Enter(const parser::CUFKernelDoConstruct &x) {
       std::get<std::list<parser::CUFReduction>>(directive.t)) {
     CheckReduce(context_, reduce);
   }
+  inCUFKernelDoConstruct_ = true;
+}
+
+void CUDAChecker::Leave(const parser::CUFKernelDoConstruct &) {
+  inCUFKernelDoConstruct_ = false;
 }
 
 void CUDAChecker::Enter(const parser::AssignmentStmt &x) {
   auto lhsLoc{std::get<parser::Variable>(x.t).GetSource()};
   const auto &scope{context_.FindScope(lhsLoc)};
   const Scope &progUnit{GetProgramUnitContaining(scope)};
-  if (IsCUDADeviceContext(&progUnit)) {
+  if (IsCUDADeviceContext(&progUnit) || inCUFKernelDoConstruct_) {
     return; // Data transfer with assignment is only perform on host.
   }
 
diff --git a/flang/lib/Semantics/check-cuda.h b/flang/lib/Semantics/check-cuda.h
index aa0cb46360bef5..222a2ee04b57c9 100644
--- a/flang/lib/Semantics/check-cuda.h
+++ b/flang/lib/Semantics/check-cuda.h
@@ -39,10 +39,12 @@ class CUDAChecker : public virtual BaseChecker {
   void Enter(const parser::FunctionSubprogram &);
   void Enter(const parser::SeparateModuleSubprogram &);
   void Enter(const parser::CUFKernelDoConstruct &);
+  void Leave(const parser::CUFKernelDoConstruct &);
   void Enter(const parser::AssignmentStmt &);
 
 private:
   SemanticsContext &context_;
+  bool inCUFKernelDoConstruct_ = false;
 };
 
 bool CanonicalizeCUDA(parser::Program &);
diff --git a/flang/test/Lower/CUDA/cuda-data-transfer.cuf b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
index f189bf9b621082..bcbac9c3031423 100644
--- a/flang/test/Lower/CUDA/cuda-data-transfer.cuf
+++ b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
@@ -332,3 +332,24 @@ end subroutine
 ! CHECK: acc.serial
 ! CHECK-NOT: cuf.data_transfer
 ! CHECK: hlfir.assign
+
+! Check that cuf.data_transfer are not generated within cuf kernel and do not
+! trigger semantic error.
+subroutine sub17()
+  integer, parameter :: n = 10
+  real, device :: adev(n)
+  real, device :: bdev(n)
+  real :: ahost
+  real, managed :: b
+  integer :: i
+
+  adev = ahost
+  !$cuf kernel do <<<*,*>>>
+  do i = 1, n
+    ahost = adev(i) * bdev(i) + b
+  enddo
+end subroutine
+
+! CHECK-LABEL: func.func @_QPsub17()
+! CHECK: cuf.kernel<<<*, *>>>
+! CHECK-NOT: cuf.data_transfer