diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h index 8cfc6adfcc941..71544a3ba5597 100644 --- a/flang/include/flang/Evaluate/tools.h +++ b/flang/include/flang/Evaluate/tools.h @@ -22,6 +22,7 @@ #include "flang/Semantics/attr.h" #include "flang/Semantics/scope.h" #include "flang/Semantics/symbol.h" +#include #include #include #include @@ -1353,6 +1354,64 @@ inline bool IsCUDAUnifiedSymbol(const Symbol &sym) { return IsCUDADataAttrSymbol(sym, common::CUDADataAttr::Unified); } +inline bool HasCUDADataAttr(const Symbol &sym) { + const auto *details{ + sym.GetUltimate().detailsIf()}; + return details && details->cudaDataAttr().has_value(); +} + +// The data attribute of a component describes the data that the component +// designates, so it hides the attribute of the object that the component is +// taken from: in a%b, where a is managed and b is device, a%b designates +// device data. Collect the symbols of the expression, leaving out the ones +// that a component with an attribute hides. +template +semantics::UnorderedSymbolSet CollectEffectiveCudaSymbols(const A &expr) { + semantics::UnorderedSymbolSet result{CollectCudaSymbols(expr)}; + SymbolVector symbols{GetSymbolVector(expr)}; + // GetSymbolVector lists the base of a component chain before its components. + // Reverse it to visit the innermost component of a chain first. + std::reverse(symbols.begin(), symbols.end()); + bool hidden{false}; + for (const Symbol &sym : symbols) { + bool isComponent{sym.owner().IsDerivedType()}; + if (hidden) { + result.erase(sym); + } else if (isComponent && HasCUDADataAttr(sym)) { + hidden = true; + } + if (!isComponent) { + hidden = false; // The base ends the component chain. + } + } + return result; +} + +// Get the number of symbols with the CUDA managed attribute in a set. +inline int CountCUDAManagedSymbols( + const semantics::UnorderedSymbolSet &symbols) { + int count{0}; + for (const Symbol &sym : symbols) { + if (IsCUDAManagedSymbol(sym)) { + ++count; + } + } + return count; +} + +// Get the number of symbols with a CUDA device attribute other than unified in +// a set. +inline int CountCUDANonUnifiedSymbols( + const semantics::UnorderedSymbolSet &symbols) { + int count{0}; + for (const Symbol &sym : symbols) { + if (IsCUDADeviceSymbol(sym) && !IsCUDAUnifiedSymbol(sym)) { + ++count; + } + } + return count; +} + // 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 @@ -1403,48 +1462,12 @@ inline int GetNbOfCUDAManagedOrUnifiedSymbols(const A &expr) { return symbols.size(); } -// Get the number of distinct symbols with the CUDA managed attribute in the -// expression. -template 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 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 inline bool HasCUDADeviceAttrs(const A &expr) { return GetNbOfCUDADeviceSymbols(expr) > 0; } -// Check if any of the symbols part of the expression has the CUDA managed -// attribute. -template inline bool HasCUDAManagedSymbols(const A &expr) { - return GetNbOfCUDAManagedSymbols(expr) > 0; -} - -// Check if any of the symbols part of the expression has a CUDA device -// attribute other than unified. -template inline bool HasCUDANonUnifiedSymbols(const A &expr) { - return GetNbOfCUDANonUnifiedSymbols(expr) > 0; -} - // 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 @@ -1470,15 +1493,20 @@ template inline bool IsWholeManagedArray(const A &expr) { // 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. +// The side of an assignment is classified from the data it designates, so the +// attribute of a component prevails over the attribute of the object it is +// taken from. // Return true if the assignment is one of the copies above. template inline bool IsCUDADataTransfer(const A &lhs, const B &rhs) { + semantics::UnorderedSymbolSet lhsSymbols{CollectEffectiveCudaSymbols(lhs)}; + semantics::UnorderedSymbolSet rhsSymbols{CollectEffectiveCudaSymbols(rhs)}; // Unified data is left out of these counts and checks so that it is handled // as host data. - bool lhsHasManaged{HasCUDAManagedSymbols(lhs)}; - bool lhsIsHost{!HasCUDANonUnifiedSymbols(lhs)}; - int rhsNbManagedSymbols{GetNbOfCUDAManagedSymbols(rhs)}; - int rhsNbSymbols{GetNbOfCUDANonUnifiedSymbols(rhs)}; + bool lhsHasManaged{CountCUDAManagedSymbols(lhsSymbols) > 0}; + bool lhsIsHost{CountCUDANonUnifiedSymbols(lhsSymbols) == 0}; + int rhsNbManagedSymbols{CountCUDAManagedSymbols(rhsSymbols)}; + int rhsNbSymbols{CountCUDANonUnifiedSymbols(rhsSymbols)}; if (HasNonAllocatableModuleCUDAManagedSymbols(lhs)) return false; diff --git a/flang/test/Lower/CUDA/cuda-managed-assign.cuf b/flang/test/Lower/CUDA/cuda-managed-assign.cuf index 09bfcd5296c89..2a201f56dedce 100644 --- a/flang/test/Lower/CUDA/cuda-managed-assign.cuf +++ b/flang/test/Lower/CUDA/cuda-managed-assign.cuf @@ -64,6 +64,45 @@ end subroutine ! CHECK-LABEL: func.func @_QPmanaged_component_attr_assign() ! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer +! The attribute of a component describes the data the component designates, so it +! prevails over the attribute of the object the component is taken from. A device +! component of a managed object is device data, which the host cannot write in +! place, whatever the form of the reference. +subroutine device_component_of_managed(n, i) + integer :: n, i + type :: t + integer(4), allocatable, device :: dcomp(:) + end type + type(t), allocatable, managed :: cm(:) + integer(4) :: ha(3) + allocate(cm(n)) + allocate(cm(i)%dcomp(3)) + cm(i)%dcomp = 10 ! device component = constant + cm(i)%dcomp = ha ! device component = host array + cm(i)%dcomp(1:2) = 10 ! device component section = constant + ha = cm(i)%dcomp ! host array = device component +end subroutine + +! CHECK-LABEL: func.func @_QPdevice_component_of_managed +! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer +! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer +! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer +! CHECK: cuf.data_transfer {{.*}}#cuf.cuda_transfer + +! The object's attribute still applies to a component that has none. +subroutine host_component_of_managed(n, i) + integer :: n, i + type :: t + integer(4) :: m(16) + end type + type(t), allocatable, managed :: cm(:) + allocate(cm(n)) + cm(i)%m = 10 +end subroutine + +! CHECK-LABEL: func.func @_QPhost_component_of_managed +! CHECK-NOT: cuf.data_transfer + ! One whole managed side is enough: the other may be a section. subroutine managed_mixed_whole_section() integer(4), managed :: ma(4), mb(4,16)