[flang-commits] [flang] [llvm] [flang][cuda] Do not use double descriptor for unified allocatables (PR #190713)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Zhen Wang (wangzpgi)

<details>
<summary>Changes</summary>

Unified memory relies on HMM/ATS and lives in host system memory, so it does not need device-side descriptor synchronization via cudaGetSymbolAddress. Treat it the same as pinned in hasDoubleDescriptor to avoid a cudaErrorInvalidSymbol at runtime.

---
Full diff: https://github.com/llvm/llvm-project/pull/190713.diff


2 Files Affected:

- (modified) flang/lib/Lower/CUDA.cpp (+2-1) 
- (modified) flang/test/Lower/CUDA/cuda-allocatable.cuf (+10) 


``````````diff
diff --git a/flang/lib/Lower/CUDA.cpp b/flang/lib/Lower/CUDA.cpp
index d8e2d829f9adf..fc7261f9a01e5 100644
--- a/flang/lib/Lower/CUDA.cpp
+++ b/flang/lib/Lower/CUDA.cpp
@@ -120,7 +120,8 @@ bool Fortran::lower::hasDoubleDescriptor(mlir::Value addr) {
     if (mlir::isa_and_nonnull<fir::AddrOfOp>(
             declareOp.getMemref().getDefiningOp())) {
       if (declareOp.getDataAttr() &&
-          *declareOp.getDataAttr() == cuf::DataAttribute::Pinned)
+          (*declareOp.getDataAttr() == cuf::DataAttribute::Pinned ||
+           *declareOp.getDataAttr() == cuf::DataAttribute::Unified))
         return false;
       return true;
     }
diff --git a/flang/test/Lower/CUDA/cuda-allocatable.cuf b/flang/test/Lower/CUDA/cuda-allocatable.cuf
index 52303d126b8dc..39179c90a2418 100644
--- a/flang/test/Lower/CUDA/cuda-allocatable.cuf
+++ b/flang/test/Lower/CUDA/cuda-allocatable.cuf
@@ -262,6 +262,16 @@ end subroutine
 ! CHECK: cuf.deallocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<managed>, hasDoubleDescriptor} -> i32
 ! CHECK: cuf.deallocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<pinned>} -> i32
 
+subroutine unified_allocate()
+  integer, allocatable, unified, save :: a_unified(:)
+  allocate(a_unified(10))
+  deallocate(a_unified)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPunified_allocate()
+! CHECK: cuf.allocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {data_attr = #cuf.cuda<unified>} -> i32
+! CHECK: cuf.deallocate %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {data_attr = #cuf.cuda<unified>} -> i32
+
 attributes(global) subroutine from_device_source()
   real, device, allocatable :: a(:)
   real, allocatable :: b(:)

``````````

</details>


https://github.com/llvm/llvm-project/pull/190713