[flang-commits] [flang] [flang][cuda] Treat unified data as host data in assignments (PR #212912)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Wed Jul 29 18:49:27 PDT 2026
https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/212912
>From 42ee13e77c3162c664e83313884e511463d2d459 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 29 Jul 2026 18:38:09 -0700
Subject: [PATCH 1/2] Treat unified data as host data in assignments
---
flang/include/flang/Evaluate/tools.h | 71 +++++--
flang/test/Lower/CUDA/cuda-unified-assign.cuf | 176 ++++++++++++++++++
2 files changed, 235 insertions(+), 12 deletions(-)
create mode 100644 flang/test/Lower/CUDA/cuda-unified-assign.cuf
diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 7d4eadbc46cb5..d27a7cec2d43b 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1336,6 +1336,22 @@ inline bool IsCUDAManagedOrUnifiedSymbol(const Symbol &sym) {
return false;
}
+inline bool IsCUDADataAttrSymbol(const Symbol &sym, common::CUDADataAttr attr) {
+ if (const auto *details =
+ sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {
+ return details->cudaDataAttr() && *details->cudaDataAttr() == attr;
+ }
+ return false;
+}
+
+inline bool IsCUDAManagedSymbol(const Symbol &sym) {
+ return IsCUDADataAttrSymbol(sym, common::CUDADataAttr::Managed);
+}
+
+inline bool IsCUDAUnifiedSymbol(const Symbol &sym) {
+ return IsCUDADataAttrSymbol(sym, common::CUDADataAttr::Unified);
+}
+
// Non-allocatable module-level managed/unified variables use pointer
// indirection through a companion global in __nv_managed_data__.
// Explicit data transfers (cudaMemcpy) must be avoided for these
@@ -1386,19 +1402,43 @@ inline int GetNbOfCUDAManagedOrUnifiedSymbols(const A &expr) {
return symbols.size();
}
+// Get the number of distinct symbols with the CUDA managed attribute in the
+// expression.
+template <typename A> inline int GetNbOfCUDAManagedSymbols(const A &expr) {
+ semantics::UnorderedSymbolSet symbols;
+ for (const Symbol &sym : CollectCudaSymbols(expr)) {
+ if (IsCUDAManagedSymbol(sym)) {
+ symbols.insert(sym);
+ }
+ }
+ return symbols.size();
+}
+
+// Get the number of distinct symbols with a CUDA device attribute other than
+// unified in the expression.
+template <typename A> inline int GetNbOfCUDANonUnifiedSymbols(const A &expr) {
+ semantics::UnorderedSymbolSet symbols;
+ for (const Symbol &sym : CollectCudaSymbols(expr)) {
+ if (IsCUDADeviceSymbol(sym) && !IsCUDAUnifiedSymbol(sym)) {
+ symbols.insert(sym);
+ }
+ }
+ return symbols.size();
+}
+
// Check if any of the symbols part of the expression has a CUDA device
// attribute.
template <typename A> inline bool HasCUDADeviceAttrs(const A &expr) {
return GetNbOfCUDADeviceSymbols(expr) > 0;
}
-// True for a whole reference to a managed or unified array: a whole array
-// variable, or a whole array component that itself has the attribute (a%b where
-// b is managed). An array section, an array element, a component of a managed
-// object and a computed value are all false.
+// True for a whole reference to a managed array: a whole array variable, or a
+// whole array component that itself has the managed attribute (a%b where b is
+// managed). An array section, an array element, a component of a managed object
+// and a computed value are all false.
template <typename A> inline bool IsWholeManagedArray(const A &expr) {
const Symbol *sym{UnwrapWholeSymbolOrComponentDataRef(expr)};
- return expr.Rank() > 0 && sym && IsCUDAManagedOrUnifiedSymbol(*sym);
+ return expr.Rank() > 0 && sym && IsCUDAManagedSymbol(*sym);
}
// CUDA Fortran Programming Guide 3.4.1 defines which assignments in host code
@@ -1414,12 +1454,19 @@ template <typename A> inline bool IsWholeManagedArray(const A &expr) {
// device data, in both directions.
// One difference from the guide is that a managed array section is copied when
// the other side is a whole managed array, as the reference compiler does.
+// The guide does not cover unified data. It is host memory that the device can
+// also access, so it takes the place of host data in the rules above and an
+// assignment between unified sides is host code, again as the reference
+// compiler does.
// Return true if the assignment is one of the copies above.
template <typename A, typename B>
inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
- int lhsNbManagedSymbols{GetNbOfCUDAManagedOrUnifiedSymbols(lhs)};
- int rhsNbManagedSymbols{GetNbOfCUDAManagedOrUnifiedSymbols(rhs)};
- int rhsNbSymbols{GetNbOfCUDADeviceSymbols(rhs)};
+ int lhsNbManagedSymbols{GetNbOfCUDAManagedSymbols(lhs)};
+ int rhsNbManagedSymbols{GetNbOfCUDAManagedSymbols(rhs)};
+ // Unified data is left out of these counts so that it is handled as host
+ // data.
+ int lhsNbSymbols{GetNbOfCUDANonUnifiedSymbols(lhs)};
+ int rhsNbSymbols{GetNbOfCUDANonUnifiedSymbols(rhs)};
if (HasNonAllocatableModuleCUDAManagedSymbols(lhs))
return false;
@@ -1436,7 +1483,7 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
// The host cannot reach device or constant data, unlike managed and unified
// data, so an assignment with such a side is a copy, sections included.
- bool lhsIsDeviceOnly{lhsNbManagedSymbols == 0 && HasCUDADeviceAttrs(lhs)};
+ bool lhsIsDeviceOnly{lhsNbManagedSymbols == 0 && lhsNbSymbols > 0};
// The right-hand side can be an expression, so one device operand is enough.
bool rhsHasDeviceOnly{rhsNbSymbols > rhsNbManagedSymbols};
@@ -1455,13 +1502,13 @@ inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
(lhsNbManagedSymbols >= 1 || rhsNbManagedSymbols >= 1)) ||
(lhsNbManagedSymbols >= 1 && !rhsHasDeviceOnly &&
!(wholeLhs || wholeRhs)) ||
- (!HasCUDADeviceAttrs(lhs) && rhsNbManagedSymbols >= 1 &&
- !rhsHasDeviceOnly && !wholeRhs) ||
+ (lhsNbSymbols == 0 && rhsNbManagedSymbols >= 1 && !rhsHasDeviceOnly &&
+ !wholeRhs) ||
(rhsNbManagedSymbols >= 1 && !IsVariable(rhs) && !lhsIsDeviceOnly) ||
(lhsNbManagedSymbols >= 1 && rhsNbSymbols == 0)) {
return false;
}
- return HasCUDADeviceAttrs(lhs) || rhsNbSymbols > 0;
+ return lhsNbSymbols > 0 || rhsNbSymbols > 0;
}
/// Check if the expression is a mix of host and device variables that require
diff --git a/flang/test/Lower/CUDA/cuda-unified-assign.cuf b/flang/test/Lower/CUDA/cuda-unified-assign.cuf
new file mode 100644
index 0000000000000..9fbb87dbae38c
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-unified-assign.cuf
@@ -0,0 +1,176 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+! Unified data is host memory that the device can also access. CUDA Fortran
+! Programming Guide 3.4.1 does not cover it, so it takes the place of host data
+! in the rules there: an assignment is a copy when the other side is device or
+! constant data, or a whole managed array, and is done on the host otherwise.
+
+module cmod
+ integer(4), constant :: ca(16)
+ integer(4), managed :: gma(16)
+end module
+
+! Assignments that are data transfers.
+
+subroutine unified_device_assign()
+ integer(4), unified :: ua(16)
+ integer(4), device :: db(16)
+ ua = db ! whole unified = whole device
+ db = ua ! whole device = whole unified
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_device_assign()
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+
+! The host cannot access device memory, so sections are copied as well.
+subroutine unified_device_section_assign(i)
+ integer :: i
+ integer(4), unified :: ua(4,16)
+ integer(4), device :: db(4,16)
+ ua(:,i) = db(:,i)
+ db(:,i) = ua(:,i)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_device_section_assign
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+
+subroutine unified_constant_assign()
+ use cmod
+ integer(4), unified :: ua(16)
+ ua = ca
+ ca = ua
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_constant_assign()
+! CHECK: cuf.data_transfer
+! CHECK: cuf.data_transfer
+
+! A whole managed array is copied whatever the other side is, and unified data on
+! the other side does not change that.
+subroutine unified_managed_assign(i)
+ integer :: i
+ integer(4), managed :: ma(4)
+ integer(4), unified :: ua(4,16), ub(4)
+ ma = ub ! whole managed = whole unified
+ ub = ma ! whole unified = whole managed
+ ua(:,i) = ma ! unified section = whole managed
+ ma = ua(:,i) ! whole managed = unified section
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_managed_assign
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+
+subroutine unified_from_managed_module()
+ use cmod
+ integer(4), unified :: ua(16)
+ ua = gma
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_from_managed_module()
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_device>
+
+! A device allocatable component of a unified object is still device data.
+subroutine unified_device_component()
+ type :: t
+ integer(4), device, allocatable :: m(:)
+ end type
+ type(t), unified :: x
+ integer(4) :: ha(5)
+ ha(1:5) = x%m(1:5)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_device_component()
+! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer<device_host>
+
+! Assignments done on the host.
+
+subroutine unified_assign(hs)
+ integer(4) :: hs
+ integer(4), unified :: ua(16), ub(16), us
+ integer(4) :: ha(16)
+ ua = ub ! unified = unified
+ ua = ha ! unified = host array
+ ha = ub ! host = unified
+ ua = 0 ! unified = constant
+ ua = hs ! unified = host scalar
+ us = hs ! unified scalar = host scalar
+ hs = us ! host scalar = unified scalar
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_assign
+! CHECK-NOT: cuf.data_transfer
+
+subroutine unified_section_assign(i)
+ integer :: i
+ integer(4), unified :: ua(4,16), ub(4,16), uc(4)
+ integer(4), managed :: ma(4,16)
+ integer(4) :: ha(4,16)
+ ua(:,i) = ub(:,i) ! unified section = unified section
+ ha(:,i) = ub(:,i) ! host section = unified section
+ ua(:,i) = ha(:,i) ! unified section = host section
+ ua(:,i) = uc ! unified section = whole unified
+ uc = ub(:,i) ! whole unified = unified section
+ ua(:,i) = ma(:,i) ! unified section = managed section
+ ma(:,i) = ub(:,i) ! managed section = unified section
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_section_assign
+! CHECK-NOT: cuf.data_transfer
+
+subroutine unified_element_assign(i)
+ integer :: i
+ integer(4), unified :: ua(16), ub(16)
+ ua(i) = ub(i)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_element_assign
+! CHECK-NOT: cuf.data_transfer
+
+! An expression is evaluated on the host, managed operands included.
+subroutine unified_expr_assign()
+ integer(4), unified :: ua(16), ub(16)
+ integer(4), managed :: mb(16)
+ integer(4) :: ha(16)
+ ua = ub + 1
+ ha = ub + 1
+ ua = mb + 1
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_expr_assign()
+! CHECK-NOT: cuf.data_transfer
+
+subroutine unified_alloc_assign(n)
+ integer :: n
+ integer(4), allocatable, unified :: ua(:)
+ integer(4), unified :: ub(16)
+ integer(4) :: ha(16)
+ allocate(ua(n))
+ ua = ub
+ ua = ha
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_alloc_assign
+! CHECK-NOT: cuf.data_transfer
+
+! The unified attribute on a component, or on the object a component belongs to,
+! does not make the assignment a copy.
+subroutine unified_component_assign()
+ type :: t
+ integer(4), unified :: m(16)
+ end type
+ type :: t2
+ integer(4) :: m(16)
+ end type
+ type(t) :: x, y
+ type(t2), unified :: ux, uy
+ x%m = y%m
+ ux%m = uy%m
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_component_assign()
+! CHECK-NOT: cuf.data_transfer
>From 4d27dcca6bf2322ab08ceab22f562d6ca0184b3f Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 29 Jul 2026 18:49:01 -0700
Subject: [PATCH 2/2] edit comments
---
flang/include/flang/Evaluate/tools.h | 7 +++----
flang/test/Lower/CUDA/cuda-unified-assign.cuf | 8 ++++----
2 files changed, 7 insertions(+), 8 deletions(-)
diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index d27a7cec2d43b..3bf4147805e7e 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1454,10 +1454,9 @@ template <typename A> inline bool IsWholeManagedArray(const A &expr) {
// device data, in both directions.
// One difference from the guide is that a managed array section is copied when
// the other side is a whole managed array, as the reference compiler does.
-// The guide does not cover unified data. It is host memory that the device can
-// also access, so it takes the place of host data in the rules above and an
-// assignment between unified sides is host code, again as the reference
-// compiler does.
+// Unified data is host memory that the device can also access, so it takes the
+// place of host data in the rules above and an assignment between unified sides
+// is host code.
// Return true if the assignment is one of the copies above.
template <typename A, typename B>
inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) {
diff --git a/flang/test/Lower/CUDA/cuda-unified-assign.cuf b/flang/test/Lower/CUDA/cuda-unified-assign.cuf
index 9fbb87dbae38c..bc4fb35d522b4 100644
--- a/flang/test/Lower/CUDA/cuda-unified-assign.cuf
+++ b/flang/test/Lower/CUDA/cuda-unified-assign.cuf
@@ -1,9 +1,9 @@
! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
-! Unified data is host memory that the device can also access. CUDA Fortran
-! Programming Guide 3.4.1 does not cover it, so it takes the place of host data
-! in the rules there: an assignment is a copy when the other side is device or
-! constant data, or a whole managed array, and is done on the host otherwise.
+! Unified data is host memory that the device can also access, so it takes the
+! place of host data in the CUDA Fortran Programming Guide 3.4.1 rules: an
+! assignment is a copy when the other side is device or constant data, or a whole
+! managed array, and is done on the host otherwise.
module cmod
integer(4), constant :: ca(16)
More information about the flang-commits
mailing list