[flang][cuda] Relax semanctic check in cuf kernel and openacc compute constructs (llvm#125750)

Previous patch was too restrictive and didn't take into account cuf
kernels and openacc compute constructs as being device context.

Author: clementval

flang/lib/Semantics/assignment.cpp

@@ -42,6 +42,7 @@ class AssignmentContext {
   void Analyze(const parser::AssignmentStmt &);
   void Analyze(const parser::PointerAssignmentStmt &);
   void Analyze(const parser::ConcurrentControl &);
+  int deviceConstructDepth_{0};
 
 private:
   bool CheckForPureContext(const SomeExpr &rhs, parser::CharBlock rhsSource);
@@ -94,7 +95,7 @@ void AssignmentContext::Analyze(const parser::AssignmentStmt &stmt) {
             common::LanguageFeature::CUDA)) {
       const auto &scope{context_.FindScope(lhsLoc)};
       const Scope &progUnit{GetProgramUnitContaining(scope)};
-      if (!IsCUDADeviceContext(&progUnit)) {
+      if (!IsCUDADeviceContext(&progUnit) && deviceConstructDepth_ == 0) {
         if (Fortran::evaluate::HasCUDADeviceAttrs(lhs) &&
             Fortran::evaluate::HasCUDAImplicitTransfer(rhs)) {
           context_.Say(lhsLoc, "Unsupported CUDA data transfer"_err_en_US);
@@ -228,6 +229,46 @@ void AssignmentChecker::Enter(const parser::MaskedElsewhereStmt &x) {
 void AssignmentChecker::Leave(const parser::MaskedElsewhereStmt &) {
   context_.value().PopWhereContext();
 }
+void AssignmentChecker::Enter(const parser::CUFKernelDoConstruct &x) {
+  ++context_.value().deviceConstructDepth_;
+}
+void AssignmentChecker::Leave(const parser::CUFKernelDoConstruct &) {
+  --context_.value().deviceConstructDepth_;
+}
+static bool IsOpenACCComputeConstruct(const parser::OpenACCBlockConstruct &x) {
+  const auto &beginBlockDirective =
+      std::get<Fortran::parser::AccBeginBlockDirective>(x.t);
+  const auto &blockDirective =
+      std::get<Fortran::parser::AccBlockDirective>(beginBlockDirective.t);
+  if (blockDirective.v == llvm::acc::ACCD_parallel ||
+      blockDirective.v == llvm::acc::ACCD_serial ||
+      blockDirective.v == llvm::acc::ACCD_kernels) {
+    return true;
+  }
+  return false;
+}
+void AssignmentChecker::Enter(const parser::OpenACCBlockConstruct &x) {
+  if (IsOpenACCComputeConstruct(x)) {
+    ++context_.value().deviceConstructDepth_;
+  }
+}
+void AssignmentChecker::Leave(const parser::OpenACCBlockConstruct &x) {
+  if (IsOpenACCComputeConstruct(x)) {
+    --context_.value().deviceConstructDepth_;
+  }
+}
+void AssignmentChecker::Enter(const parser::OpenACCCombinedConstruct &) {
+  ++context_.value().deviceConstructDepth_;
+}
+void AssignmentChecker::Leave(const parser::OpenACCCombinedConstruct &) {
+  --context_.value().deviceConstructDepth_;
+}
+void AssignmentChecker::Enter(const parser::OpenACCLoopConstruct &) {
+  ++context_.value().deviceConstructDepth_;
+}
+void AssignmentChecker::Leave(const parser::OpenACCLoopConstruct &) {
+  --context_.value().deviceConstructDepth_;
+}
 
 } // namespace Fortran::semantics
 template class Fortran::common::Indirection<

flang/lib/Semantics/assignment.h

@@ -45,6 +45,14 @@ class AssignmentChecker : public virtual BaseChecker {
   void Leave(const parser::EndWhereStmt &);
   void Enter(const parser::MaskedElsewhereStmt &);
   void Leave(const parser::MaskedElsewhereStmt &);
+  void Enter(const parser::CUFKernelDoConstruct &);
+  void Leave(const parser::CUFKernelDoConstruct &);
+  void Enter(const parser::OpenACCBlockConstruct &);
+  void Leave(const parser::OpenACCBlockConstruct &);
+  void Enter(const parser::OpenACCCombinedConstruct &);
+  void Leave(const parser::OpenACCCombinedConstruct &);
+  void Enter(const parser::OpenACCLoopConstruct &);
+  void Leave(const parser::OpenACCLoopConstruct &);
 
 private:
   common::Indirection<AssignmentContext> context_;

flang/test/Semantics/cuf18.cuf

@@ -1,10 +1,67 @@
-! RUN: %python %S/test_errors.py %s %flang_fc1
+! RUN: %python %S/test_errors.py %s %flang_fc1 -fopenacc
 
 subroutine sub1()
   real, allocatable, device :: a(:)
+  integer :: i
 
 !ERROR: Unsupported CUDA data transfer
   a = a + 10 ! Illegal expression according to 3.4.2
+
+  !$cuf kernel do
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in cuf kernel do
+  end do
+
+  !$acc parallel loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in openacc combined construct
+  end do
+
+  !$acc serial loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in openacc combined construct
+  end do
+
+  !$acc kernels loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in openacc combined construct
+  end do
+
+  !$acc parallel
+  !$acc loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in nested openacc construct
+  end do
+  !$acc end parallel
+
+  !$acc kernels
+  !$acc loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in nested openacc construct
+  end do
+  !$acc end kernels
+
+  !$acc serial
+  !$acc loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok in nested openacc construct
+  end do
+  !$acc end serial
+
+  !$acc loop
+  do i = 1, 10
+    a(i) = a(i) + 10 ! ok acc loop
+  end do
+
+  !$acc data
+
+  do i = 1, 10
+!ERROR: Unsupported CUDA data transfer
+    a(i) = a(i) + 10
+  end do
+
+  !$acc end data
+
 end subroutine