[flang-commits] [flang] 7a3208c - [flang][cuda] Emit data transfers for whole-array and function-result managed assignments (#211640)

via flang-commits flang-commits at lists.llvm.org
Fri Jul 24 08:41:09 PDT 2026


Author: Zhen Wang
Date: 2026-07-24T15:41:03Z
New Revision: 7a3208c8550564489037111f1d85a01f18a196a2

URL: https://github.com/llvm/llvm-project/commit/7a3208c8550564489037111f1d85a01f18a196a2
DIFF: https://github.com/llvm/llvm-project/commit/7a3208c8550564489037111f1d85a01f18a196a2.diff

LOG: [flang][cuda] Emit data transfers for whole-array and function-result managed assignments (#211640)

Assignments that read or write device/managed/unified data are
synchronous data transfers on stream 0 that wait for previously launched
kernels. Previously, whole-array managed = managed / host = managed
assignments, and assignments consuming a managed/unified/device function
result, were lowered as host copies (hlfir.assign), which can race an
asynchronous kernel that produced the data.

Emit cuf.data_transfer for these cases. Element-wise (scalar) accesses
and right-hand-side expressions involving managed data remain host
operations (the latter avoids materializing a temporary), and a
whole-allocatable left-hand side keeps reallocation semantics.

Added: 
    flang/test/Lower/CUDA/cuda-managed-assign.cuf

Modified: 
    flang/include/flang/Evaluate/tools.h
    flang/lib/Lower/Bridge.cpp
    flang/test/Lower/CUDA/cuda-data-transfer.cuf

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 6d602da2a3e80..45c45b78bb694 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1408,13 +1408,27 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
     return true; // Managed arrays initialization is performed on the device.
   }
 
-  // Cases where no explicit data transfer is needed:
-  // - Both sides involve only managed/unified symbols (host-accessible).
-  // - LHS is host-only and RHS has only managed/unified symbols.
-  // - LHS is managed/unified and RHS is host-only.
-  if ((lhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols) ||
+  // Managed/unified data is host-addressable, so several assignments are
+  // performed on the host and need no explicit data transfer:
+  // - A whole-allocatable left-hand side involving managed/unified data: the
+  //   assignment has reallocation semantics and is performed on the host.
+  // - Element-wise (scalar) access to managed/unified data.
+  // - A right-hand side expression involving managed/unified data assigned into
+  //   a host-addressable (managed/unified or host) left-hand side: evaluating
+  //   it on the host avoids materializing a temporary.
+  // - A managed/unified left-hand side assigned from host-only data.
+  // A whole-array assignment whose right-hand side is a managed/unified
+  // variable is a synchronous data transfer that waits for previously launched
+  // kernels.
+  if ((IsAllocatableDesignator(lhs) &&
+          (lhsNbManagedSymbols >= 1 || rhsNbManagedSymbols >= 1)) ||
+      (lhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols &&
+          lhs.Rank() == 0) ||
       (lhsNbManagedSymbols == 0 && !HasCUDADeviceAttrs(lhs) &&
-          rhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols) ||
+          rhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols &&
+          lhs.Rank() == 0) ||
+      (rhsNbManagedSymbols >= 1 && !IsVariable(rhs) &&
+          (lhsNbManagedSymbols >= 1 || !HasCUDADeviceAttrs(lhs))) ||
       (lhsNbManagedSymbols >= 1 && rhsNbSymbols == 0)) {
     return false;
   }

diff  --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 7cde3ae492e14..fec9106a9ec38 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -5270,13 +5270,44 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     return false;
   }
 
+  // Return true if the right-hand side of the assignment is a reference to a
+  // function whose result carries a managed, unified, or device CUDA data
+  // attribute. Such a result may be produced by an asynchronous kernel, so
+  // consuming it in an assignment must be a synchronizing data transfer rather
+  // than a plain host assignment. A whole-allocatable left-hand side is
+  // excluded: it has reallocation semantics and is performed on the host.
+  bool
+  isCUDAFunctionResultTransfer(const Fortran::evaluate::Assignment &assign) {
+    if (Fortran::evaluate::IsAllocatableDesignator(assign.lhs))
+      return false;
+    const Fortran::evaluate::ProcedureRef *procRef =
+        Fortran::evaluate::UnwrapProcedureRef(assign.rhs);
+    if (!procRef)
+      return false;
+    auto procedure =
+        Fortran::evaluate::characteristics::Procedure::Characterize(
+            procRef->proc(), getFoldingContext(), /*emitError=*/false);
+    if (!procedure || !procedure->functionResult ||
+        !procedure->functionResult->cudaDataAttr)
+      return false;
+    Fortran::common::CUDADataAttr attr =
+        *procedure->functionResult->cudaDataAttr;
+    return attr == Fortran::common::CUDADataAttr::Managed ||
+           attr == Fortran::common::CUDADataAttr::Unified ||
+           attr == Fortran::common::CUDADataAttr::Device;
+  }
+
   void genCUDADataTransfer(fir::FirOpBuilder &builder, mlir::Location loc,
                            const Fortran::evaluate::Assignment &assign,
                            hlfir::Entity &lhs, hlfir::Entity &rhs,
                            bool isWholeAllocatableAssignment,
                            bool keepLhsLengthInAllocatableAssignment) {
     bool lhsIsDevice = Fortran::evaluate::HasCUDADeviceAttrs(assign.lhs);
-    bool rhsIsDevice = Fortran::evaluate::HasCUDADeviceAttrs(assign.rhs);
+    // A managed/unified/device function result is not visible to the symbol
+    // collection used by HasCUDADeviceAttrs (a ProcedureRef contributes no
+    // symbols), so treat such a result as a device side for transfer direction.
+    bool rhsIsDevice = Fortran::evaluate::HasCUDADeviceAttrs(assign.rhs) ||
+                       isCUDAFunctionResultTransfer(assign);
     mlir::UnitAttr hasManagedOrUnifedSymbols =
         (Fortran::evaluate::GetNbOfCUDAManagedOrUnifiedSymbols(assign.lhs) >
              0 ||
@@ -5459,8 +5490,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         getFoldingContext().languageFeatures().IsEnabled(
             Fortran::common::LanguageFeature::DoConcurrentOffload));
 
-    bool isCUDATransfer =
-        IsCUDADataTransfer(assign.lhs, assign.rhs) && !isInDeviceContext;
+    bool isCUDATransfer = (IsCUDADataTransfer(assign.lhs, assign.rhs) ||
+                           isCUDAFunctionResultTransfer(assign)) &&
+                          !isInDeviceContext;
     bool hasCUDAImplicitTransfer =
         isCUDATransfer &&
         Fortran::evaluate::HasCUDAImplicitTransfer(assign.rhs);

diff  --git a/flang/test/Lower/CUDA/cuda-data-transfer.cuf b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
index f1cd5bae81419..439a5f4bfb77c 100644
--- a/flang/test/Lower/CUDA/cuda-data-transfer.cuf
+++ b/flang/test/Lower/CUDA/cuda-data-transfer.cuf
@@ -560,9 +560,7 @@ end subroutine
 
 ! CHECK-LABEL: func.func @_QPsub28
 ! CHECK: fir.call @_QP__sum
-! CHECK-NOT: cuf.data_transfer
-! CHECK: hlfir.assign
-! CHECK-NOT: cuf.data_transfer
+! CHECK: cuf.data_transfer
 
 ! Data transfer with conversion with more complex elemental
 ! Check that the data transfer is placed before the elemental op.
@@ -655,8 +653,8 @@ end subroutine
 ! CHECK-LABEL: func.func @_QPsub35()
 ! CHECK-NOT: cuf.data_transfer
 
-! Test that host_var = managed_module_var does NOT generate cuf.data_transfer
-! (managed memory is host-accessible, so direct assignment suffices).
+! Test that host_var = managed_module_var generates a cuf.data_transfer:
+! reading a whole managed array variable is a synchronous data transfer.
 subroutine sub36()
   use managed_mod
   integer :: host_arr(10)
@@ -664,7 +662,7 @@ subroutine sub36()
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPsub36()
-! CHECK-NOT: cuf.data_transfer
+! CHECK: cuf.data_transfer %{{.*}} to %{{.*}} {hasManagedOrUnifedSymbols, transfer_kind = #cuf.cuda_transfer<device_host>}
 
 ! Test that device_var = managed_module_var DOES generate cuf.data_transfer
 ! (device memory requires explicit cudaMemcpy).

diff  --git a/flang/test/Lower/CUDA/cuda-managed-assign.cuf b/flang/test/Lower/CUDA/cuda-managed-assign.cuf
new file mode 100644
index 0000000000000..126169d3d40b3
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-managed-assign.cuf
@@ -0,0 +1,65 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! A whole-array assignment whose right-hand side is a managed variable or a
+! managed function result is a synchronous data transfer (matching CUDA Fortran
+! assignment-statement semantics). Element-wise (scalar) accesses and right-hand
+! side expressions involving managed data are performed on the host and need no
+! transfer.
+
+module mfr
+contains
+  function fr(n) result(r)
+    integer, value :: n
+    integer(4), allocatable, managed :: r(:)
+    allocate(r(n))
+    r = 7
+  end function
+end module
+
+subroutine managed_array_assign()
+  integer(4), managed :: ma(16), mb(16)
+  integer(4) :: ha(16)
+  ma = mb    ! managed = managed
+  ma = ha    ! managed = host
+  ha = ma    ! host = managed
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_array_assign()
+! CHECK: cuf.data_transfer %{{.*}} to %{{.*}} {hasManagedOrUnifedSymbols, transfer_kind = #cuf.cuda_transfer<device_device>}
+! CHECK: cuf.data_transfer %{{.*}} to %{{.*}} {hasManagedOrUnifedSymbols, transfer_kind = #cuf.cuda_transfer<host_device>}
+! CHECK: cuf.data_transfer %{{.*}} to %{{.*}} {hasManagedOrUnifedSymbols, transfer_kind = #cuf.cuda_transfer<device_host>}
+
+! A right-hand side expression involving managed data is evaluated on the host
+! to avoid materializing a temporary, so no data transfer is generated.
+subroutine managed_expr_assign()
+  integer(4), managed :: ma(16), mb(16)
+  integer(4) :: ha(16)
+  ma = mb + 1    ! managed = managed expression
+  ha = ma + 1    ! host = managed expression
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_expr_assign()
+! CHECK-NOT: cuf.data_transfer
+
+subroutine managed_scalar_access(n)
+  integer :: n, i
+  integer(4), managed :: ma(16)
+  integer(4) :: ha(16)
+  do i = 1, n
+    ha(i) = ma(i) + 1   ! element-wise access is performed on the host
+  end do
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_scalar_access
+! CHECK-NOT: cuf.data_transfer
+
+! A managed function result may be produced by an asynchronous kernel, so
+! consuming it in an assignment is a synchronizing data transfer.
+subroutine managed_func_result()
+  use mfr
+  integer(4), managed :: b(4)
+  b = fr(4)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmanaged_func_result()
+! CHECK: cuf.data_transfer


        


More information about the flang-commits mailing list