[flang][cuda] Add restriction on implicit data transfer

[clementval]

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. This patch adds a similar restriction.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

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

Changes

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.

---

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

3 Files Affected:

• (modified) flang/include/flang/Evaluate/tools.h (+11-5)
• (modified) flang/lib/Semantics/check-cuda.cpp (+16)
• (modified) flang/lib/Semantics/check-cuda.h (+2)

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_;