[flang-commits] [flang] [flang][cuda] Emit data transfers for whole-array and function-result managed assignments (PR #211640)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Fri Jul 24 08:12:28 PDT 2026
https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/211640
>From e111e254c33b2dbc9d28eef4c9e7853204cbbfe6 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 22 Jul 2026 15:29:54 -0700
Subject: [PATCH 1/3] Emit cuf.data_transfer for whole-array managed=managed
and host=managed assignments
---
flang/include/flang/Evaluate/tools.h | 21 ++++++---
flang/test/Lower/CUDA/cuda-managed-assign.cuf | 43 +++++++++++++++++++
2 files changed, 58 insertions(+), 6 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-managed-assign.cuf
diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 6d602da2a3e80..7cb895406146f 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1408,13 +1408,22 @@ 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:
+ // - 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 ((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/test/Lower/CUDA/cuda-managed-assign.cuf b/flang/test/Lower/CUDA/cuda-managed-assign.cuf
new file mode 100644
index 0000000000000..b05e8b37030f2
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-managed-assign.cuf
@@ -0,0 +1,43 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! A whole-array assignment whose right-hand side is a managed variable 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.
+
+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
>From 00203363114c8b30f6af6237b9b5f0109a99068d Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 23 Jul 2026 09:47:42 -0700
Subject: [PATCH 2/3] func result assginment
---
flang/include/flang/Evaluate/tools.h | 6 +++-
flang/lib/Lower/Bridge.cpp | 36 +++++++++++++++++--
flang/test/Lower/CUDA/cuda-data-transfer.cuf | 10 +++---
flang/test/Lower/CUDA/cuda-managed-assign.cuf | 30 +++++++++++++---
4 files changed, 69 insertions(+), 13 deletions(-)
diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 7cb895406146f..a05c7fa07676e 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1410,6 +1410,8 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
// 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
@@ -1417,7 +1419,9 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
// - 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 ((lhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols &&
+ if ((IsAllocatableDesignator(lhs) &&
+ (lhsNbManagedSymbols >= 1 || rhsNbManagedSymbols >= 1)) ||
+ (lhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols &&
lhs.Rank() == 0) ||
(lhsNbManagedSymbols == 0 && !HasCUDADeviceAttrs(lhs) &&
rhsNbManagedSymbols >= 1 && rhsNbManagedSymbols == rhsNbSymbols &&
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 7cde3ae492e14..6e42692356aef 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -5270,13 +5270,43 @@ 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 ||
@@ -5460,7 +5490,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
Fortran::common::LanguageFeature::DoConcurrentOffload));
bool isCUDATransfer =
- IsCUDADataTransfer(assign.lhs, assign.rhs) && !isInDeviceContext;
+ (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
index b05e8b37030f2..126169d3d40b3 100644
--- a/flang/test/Lower/CUDA/cuda-managed-assign.cuf
+++ b/flang/test/Lower/CUDA/cuda-managed-assign.cuf
@@ -1,9 +1,20 @@
! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
-! A whole-array assignment whose right-hand side is a managed variable 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.
+! 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)
@@ -41,3 +52,14 @@ 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
>From e043fcbbd024f57ffaec512469d4280237dc97e9 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 23 Jul 2026 11:45:00 -0700
Subject: [PATCH 3/3] format
---
flang/include/flang/Evaluate/tools.h | 9 +++++----
flang/lib/Lower/Bridge.cpp | 16 ++++++++--------
2 files changed, 13 insertions(+), 12 deletions(-)
diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index a05c7fa07676e..45c45b78bb694 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1414,11 +1414,12 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
// 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 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.
+ // 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 &&
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 6e42692356aef..fec9106a9ec38 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -5276,16 +5276,17 @@ class FirConverter : public Fortran::lower::AbstractConverter {
// 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) {
+ 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);
+ auto procedure =
+ Fortran::evaluate::characteristics::Procedure::Characterize(
+ procRef->proc(), getFoldingContext(), /*emitError=*/false);
if (!procedure || !procedure->functionResult ||
!procedure->functionResult->cudaDataAttr)
return false;
@@ -5489,10 +5490,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
getFoldingContext().languageFeatures().IsEnabled(
Fortran::common::LanguageFeature::DoConcurrentOffload));
- bool isCUDATransfer =
- (IsCUDADataTransfer(assign.lhs, assign.rhs) ||
- isCUDAFunctionResultTransfer(assign)) &&
- !isInDeviceContext;
+ bool isCUDATransfer = (IsCUDADataTransfer(assign.lhs, assign.rhs) ||
+ isCUDAFunctionResultTransfer(assign)) &&
+ !isInDeviceContext;
bool hasCUDAImplicitTransfer =
isCUDATransfer &&
Fortran::evaluate::HasCUDAImplicitTransfer(assign.rhs);
More information about the flang-commits
mailing list