1 files changed, 13 insertions, 4 deletions
diff --git a/flang/lib/Semantics/check-cuda.cpp b/flang/lib/Semantics/check-cuda.cpp
index 8d95a3a..3d2db6a 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;
   }
 
@@ -785,9 +786,17 @@ void CUDAChecker::Enter(const parser::PrintStmt &x) {
   for (const auto &item : outputItemList) {
     if (const auto *x{std::get_if<parser::Expr>(&item.u)}) {
       if (const auto *expr{GetExpr(context_, *x)}) {
-        if (Fortran::evaluate::HasCUDADeviceAttrs(*expr)) {
-          context_.Say(parser::FindSourceLocation(*x),
-              "device data not allowed in I/O statements"_err_en_US);
+        for (const Symbol &sym : CollectCudaSymbols(*expr)) {
+          if (const auto *details = sym.GetUltimate()
+                  .detailsIf<semantics::ObjectEntityDetails>()) {
+            if (details->cudaDataAttr() &&
+                (*details->cudaDataAttr() == common::CUDADataAttr::Device ||
+                    *details->cudaDataAttr() ==
+                        common::CUDADataAttr::Constant)) {
+              context_.Say(parser::FindSourceLocation(*x),
+                  "device data not allowed in I/O statements"_err_en_US);
+            }
+          }
         }
       }
     }