[flang-commits] [flang] [flang][cuda] Add restriction on implicit data transfer (PR #87720)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/87720

In section 3.4.2, some example of illegal data transfer using expression are given. One of it is when multiple device objects are part of an expression in the rhs. Current implementation allow a single device object in such case.


>From 74aed2fa8f97da3083a2080e09b91f7cd9f4eb0d Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 4 Apr 2024 15:05:19 -0700
Subject: [PATCH] [flang][cuda] Add restriction on implicit data transfer

In section 3.4.2, some example of illegal data transfer using expression
are given. One of it is when multiple device objects are part of an
expression in the rhs. Current implementation allow a single device object
in such case.
---
 flang/include/flang/Evaluate/tools.h | 16 +++++++++++-----
 flang/lib/Semantics/check-cuda.cpp   | 16 ++++++++++++++++
 flang/lib/Semantics/check-cuda.h     |  2 ++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 8c872a0579c8ed..75498f0f685844 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1226,18 +1226,24 @@ bool CheckForCoindexedObject(parser::ContextualMessages &,
     const std::optional<ActualArgument> &, const std::string &procName,
     const std::string &argName);
 
-/// Check if any of the symbols part of the expression has a cuda data
-/// attribute.
-inline bool HasCUDAAttrs(const Expr<SomeType> &expr) {
+// Get the number of symbols with CUDA attribute in the expression.
+template <typename A> inline unsigned GetNbOfCUDASymbols(const A &expr) {
+  unsigned n{0};
   for (const Symbol &sym : CollectSymbols(expr)) {
     if (const auto *details =
             sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {
       if (details->cudaDataAttr()) {
-        return true;
+        ++n;
       }
     }
   }
-  return false;
+  return n;
+}
+
+// Check if any of the symbols part of the expression has a CUDA data
+// attribute.
+template <typename A> inline bool HasCUDAAttrs(const A &expr) {
+  return GetNbOfCUDASymbols(expr) > 0;
 }
 
 } // namespace Fortran::evaluate
diff --git a/flang/lib/Semantics/check-cuda.cpp b/flang/lib/Semantics/check-cuda.cpp
index c0c6ff4c1a2ba3..c0b4e53f30c08c 100644
--- a/flang/lib/Semantics/check-cuda.cpp
+++ b/flang/lib/Semantics/check-cuda.cpp
@@ -9,12 +9,14 @@
 #include "check-cuda.h"
 #include "flang/Common/template.h"
 #include "flang/Evaluate/fold.h"
+#include "flang/Evaluate/tools.h"
 #include "flang/Evaluate/traverse.h"
 #include "flang/Parser/parse-tree-visitor.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Parser/tools.h"
 #include "flang/Semantics/expression.h"
 #include "flang/Semantics/symbol.h"
+#include "flang/Semantics/tools.h"
 
 // Once labeled DO constructs have been canonicalized and their parse subtrees
 // transformed into parser::DoConstructs, scan the parser::Blocks of the program
@@ -413,4 +415,18 @@ void CUDAChecker::Enter(const parser::CUFKernelDoConstruct &x) {
   }
 }
 
+void CUDAChecker::Enter(const parser::AssignmentStmt &x) {
+  const evaluate::Assignment *assign = semantics::GetAssignment(x);
+  unsigned nbLhs = evaluate::GetNbOfCUDASymbols(assign->lhs);
+  unsigned nbRhs = evaluate::GetNbOfCUDASymbols(assign->rhs);
+  auto lhsLoc{std::get<parser::Variable>(x.t).GetSource()};
+
+  // device to host transfer with more than one device object on the rhs is not
+  // legal.
+  if (nbLhs == 0 && nbRhs > 1) {
+    context_.Say(lhsLoc,
+        "More than one reference to a CUDA object on the right hand side of the assigment"_err_en_US);
+  }
+}
+
 } // namespace Fortran::semantics
diff --git a/flang/lib/Semantics/check-cuda.h b/flang/lib/Semantics/check-cuda.h
index d863795f16a7c2..aa0cb46360bef5 100644
--- a/flang/lib/Semantics/check-cuda.h
+++ b/flang/lib/Semantics/check-cuda.h
@@ -17,6 +17,7 @@ struct Program;
 class Messages;
 struct Name;
 class CharBlock;
+struct AssignmentStmt;
 struct ExecutionPartConstruct;
 struct ExecutableConstruct;
 struct ActionStmt;
@@ -38,6 +39,7 @@ class CUDAChecker : public virtual BaseChecker {
   void Enter(const parser::FunctionSubprogram &);
   void Enter(const parser::SeparateModuleSubprogram &);
   void Enter(const parser::CUFKernelDoConstruct &);
+  void Enter(const parser::AssignmentStmt &);
 
 private:
   SemanticsContext &context_;