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

[clementval]

The switch to GetSymbolVector introduced a regression on detecting implicit data transfer when the rhs is a function call.
Make sure the symbol we are looking at are of interest to detect data transfer.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

@llvm/pr-subscribers-flang-fir-hlfir

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

Changes

The switch to GetSymbolVector introduced a regression on detecting implicit data transfer when the rhs is a function call.
Make sure the symbol we are looking at are of interest to detect data transfer.

---

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

2 Files Affected:

• (modified) flang/lib/Evaluate/tools.cpp (+15-10)
• (modified) flang/test/Lower/CUDA/cuda-data-transfer.cuf (+8)

diff --git a/flang/lib/Evaluate/tools.cpp b/flang/lib/Evaluate/tools.cpp
index 3d9f06308d8c1..21e6b3c3dd50d 100644
--- a/flang/lib/Evaluate/tools.cpp
+++ b/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;
diff --git a/flang/test/Lower/CUDA/cuda-data-transfer.cuf b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
index 3a9b55996d9b1..3a4aff977b7a5 100644
--- a/flang/test/Lower/CUDA/cuda-data-transfer.cuf
+++ b/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