[flang][cuda] Add more support for data transfer with constant (#211944)

Extent support where rhs has some constant variables in an expression.
This match what is accepted in the legacy compiler.

flang/include/flang/Evaluate/tools.h

@@ -1439,6 +1439,9 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
 /// implicit data transfer.
 bool HasCUDAImplicitTransfer(const Expr<SomeType> &expr);
 
+/// Check if the expression is a mix of host and constant variables.
+bool HasOnlyCUDAConstntImplicitTransfer(const Expr<SomeType> &expr);
+
 // Checks whether the symbol on the LHS is present in the RHS expression.
 bool CheckForSymbolMatch(const Expr<SomeType> *lhs, const Expr<SomeType> *rhs);
 

flang/lib/Evaluate/tools.cpp

@@ -1246,7 +1246,8 @@ int GetNbOfUniqueCUDADeviceSymbols(const Expr<SomeType> &expr) {
   return symbols.size();
 }
 
-bool HasCUDAImplicitTransfer(const Expr<SomeType> &expr) {
+std::pair<semantics::UnorderedSymbolSet, semantics::UnorderedSymbolSet>
+GetHostAndDeviceSymbols(const Expr<SomeType> &expr) {
   semantics::UnorderedSymbolSet hostSymbols;
   semantics::UnorderedSymbolSet deviceSymbols;
   semantics::UnorderedSymbolSet cudaSymbols{CollectCudaSymbols(expr)};
@@ -1272,10 +1273,29 @@ bool HasCUDAImplicitTransfer(const Expr<SomeType> &expr) {
       skipNext = false;
     }
   }
+  return std::make_pair(hostSymbols, deviceSymbols);
+}
+
+bool HasCUDAImplicitTransfer(const Expr<SomeType> &expr) {
+  auto [hostSymbols, deviceSymbols] = GetHostAndDeviceSymbols(expr);
   bool hasConstant{HasConstant(expr)};
   return (hasConstant || (hostSymbols.size() > 0)) && deviceSymbols.size() > 0;
 }
 
+bool HasOnlyCUDAConstntImplicitTransfer(const Expr<SomeType> &expr) {
+  auto [hostSymbols, deviceSymbols] = GetHostAndDeviceSymbols(expr);
+  for (const Symbol &sym : deviceSymbols) {
+    if (const auto *details =
+            sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {
+      if (details->cudaDataAttr() &&
+          (*details->cudaDataAttr() != common::CUDADataAttr::Constant)) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
 bool IsCUDADeviceSymbol(const Symbol &sym) {
   if (const auto *details =
           sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {

flang/lib/Lower/Bridge.cpp

@@ -5590,9 +5590,22 @@ private:
           if (lhs.getDefiningOp())
             attachInlineAttributes(*lhs.getDefiningOp(), dirs);
         }
-        hlfir::AssignOp::create(builder, loc, rhs, lhs,
-                                isWholeAllocatableAssignment,
-                                keepLhsLengthInAllocatableAssignment);
+        if (isCUDATransfer && hasCUDAImplicitTransfer &&
+            Fortran::evaluate::HasCUDADeviceAttrs(assign.lhs)) {
+          auto [temp, cleanup] = hlfir::createTempFromMold(loc, builder, lhs);
+          hlfir::AssignOp::create(builder, loc, rhs, temp,
+                                  isWholeAllocatableAssignment,
+                                  keepLhsLengthInAllocatableAssignment);
+          auto transferKindAttr = cuf::DataTransferKindAttr::get(
+              builder.getContext(), cuf::DataTransferKind::HostDevice);
+          cuf::DataTransferOp::create(builder, loc, temp, lhs,
+                                      /*shape=*/mlir::Value{},
+                                      transferKindAttr);
+        } else {
+          hlfir::AssignOp::create(builder, loc, rhs, lhs,
+                                  isWholeAllocatableAssignment,
+                                  keepLhsLengthInAllocatableAssignment);
+        }
       }
       if (hasCUDAImplicitTransfer && !isInDeviceContext) {
         localSymbols.popScope();

flang/lib/Semantics/check-cuda.cpp

@@ -842,7 +842,8 @@ void CUDAChecker::Enter(const parser::AssignmentStmt &x) {
   }
 
   if (evaluate::HasCUDADeviceAttrs(assign->lhs) &&
-      evaluate::HasCUDAImplicitTransfer(assign->rhs)) {
+      (evaluate::HasCUDAImplicitTransfer(assign->rhs) &&
+          !evaluate::HasOnlyCUDAConstntImplicitTransfer(assign->rhs))) {
     if (GetNbOfCUDAManagedOrUnifiedSymbols(assign->lhs) == 1 &&
         GetNbOfCUDAManagedOrUnifiedSymbols(assign->rhs) == 1 && nbRhs == 1) {
       return; // This is a special case handled on the host.

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

@@ -3,6 +3,11 @@
 ! Test CUDA Fortran data transfer using assignment statements.
 
 module mod1
+
+  real, parameter :: pi = 0.9189385332046727417803297
+  real, constant :: rconst
+  real, constant :: csconst
+
   type :: t1
     integer :: i
   end type
@@ -774,3 +779,16 @@ end subroutine
 ! CHECK-LABEL: func.func @_QPsub45
 ! CHECK-LABEL: @_QFsub45Pcompute
 ! CHECK: cuf.data_transfer %{{.*}} to %{{.*}} {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.box<!fir.array<?x?xcomplex<f64>>>, !fir.box<!fir.array<?x?xcomplex<f64>>>
+
+subroutine sub46()
+  use mod1
+  csconst = rconst * pi/6
+end subroutine
+
+! CHECK-LABEL: func.func @_QPsub46()
+! CHECK: cuf.data_transfer
+! CHECK: arith.mulf
+! CHECK: arith.constant
+! CHECK: arith.divf
+! CHECK: hlfir.assign
+! CHECK: cuf.data_transfer