[flang][cuda] Fix false positive on unsupported CUDA data transfer

Author: clementval

flang/lib/Evaluate/tools.cpp

@@ -1157,23 +1157,28 @@ template semantics::UnorderedSymbolSet CollectCudaSymbols(
 bool HasCUDAImplicitTransfer(const Expr<SomeType> &expr) {
   semantics::UnorderedSymbolSet hostSymbols;
   semantics::UnorderedSymbolSet deviceSymbols;
+  semantics::UnorderedSymbolSet cudaSymbols{CollectCudaSymbols(expr)};
 
   SymbolVector symbols{GetSymbolVector(expr)};
   std::reverse(symbols.begin(), symbols.end());
   bool skipNext{false};
   for (const Symbol &sym : symbols) {
-    bool isComponent{sym.owner().IsDerivedType()};
-    bool skipComponent{false};
-    if (!skipNext) {
-      if (IsCUDADeviceSymbol(sym)) {
-        deviceSymbols.insert(sym);
-      } else if (isComponent) {
-        skipComponent = true; // Component is not device. Look on the base.
-      } else {
-        hostSymbols.insert(sym);
+    if (cudaSymbols.find(sym) != cudaSymbols.end()) {
+      bool isComponent{sym.owner().IsDerivedType()};
+      bool skipComponent{false};
+      if (!skipNext) {
+        if (IsCUDADeviceSymbol(sym)) {
+          deviceSymbols.insert(sym);
+        } else if (isComponent) {
+          skipComponent = true; // Component is not device. Look on the base.
+        } else {
+          hostSymbols.insert(sym);
+        }
       }
+      skipNext = isComponent && !skipComponent;
+    } else {
+      skipNext = false;
     }
-    skipNext = isComponent && !skipComponent;
   }
   bool hasConstant{HasConstant(expr)};
   return (hasConstant || (hostSymbols.size() > 0)) && deviceSymbols.size() > 0;

flang/test/Lower/CUDA/cuda-data-transfer.cuf

@@ -436,3 +436,11 @@ end subroutine
 
 ! CHECK-LABEL: func.func @_QPsub22()
 ! CHECK: cuf.data_transfer
+
+subroutine sub23(n)
+  integer :: n
+  real(8), device :: d(n,n), x(n)
+  x = sum(d,dim=2) ! Was triggering Unsupported CUDA data transfer
+end subroutine
+
+! CHECK-LABEL: func.func @_QPsub23