[flang][cuda] Avoid I/O error in block inside a kernel

[clementval]

Make sure we get the program unit to check the device context. The scope would be the block otherwise.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

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

Changes

Make sure we get the program unit to check the device context. The scope would be the block otherwise.

---

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

2 Files Affected:

• (modified) flang/lib/Semantics/check-cuda.cpp (+2-1)
• (modified) flang/test/Semantics/cuf23.cuf (+7)

diff --git a/flang/lib/Semantics/check-cuda.cpp b/flang/lib/Semantics/check-cuda.cpp
index 8d95a3ab580ef..3e5c4f6dff0b5 100644
--- a/flang/lib/Semantics/check-cuda.cpp
+++ b/flang/lib/Semantics/check-cuda.cpp
@@ -777,7 +777,8 @@ void CUDAChecker::Enter(const parser::AssignmentStmt &x) {
 void CUDAChecker::Enter(const parser::PrintStmt &x) {
   CHECK(context_.location());
   const Scope &scope{context_.FindScope(*context_.location())};
-  if (IsCUDADeviceContext(&scope) || deviceConstructDepth_ > 0) {
+  const Scope &progUnit{GetProgramUnitContaining(scope)};
+  if (IsCUDADeviceContext(&progUnit) || deviceConstructDepth_ > 0) {
     return;
   }
 
diff --git a/flang/test/Semantics/cuf23.cuf b/flang/test/Semantics/cuf23.cuf
index 386ad50a70acb..34a5d822cd586 100644
--- a/flang/test/Semantics/cuf23.cuf
+++ b/flang/test/Semantics/cuf23.cuf
@@ -32,3 +32,10 @@ attributes(device) subroutine device1()
   real, device :: a(10)
   print*, a ! ok
 end subroutine
+
+attributes(global) subroutine global_with_block()
+  block
+    real, device :: a(10)
+    print*, a ! ok
+  end block
+end subroutine