[flang-commits] [flang] [flang][cuda] Allocate adjustable automatic arrays in unified/managed memory (PR #212965)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: Matsu (khaki3)

<details>
<summary>Changes</summary>

Example:
```fortran
subroutine work(n)
  integer :: n
  complex(8) :: ar_tmp(n)   ! adjustable/VLA automatic
  !$acc host_data use_device(ar_tmp)
  ! ... device use ...
  !$acc end host_data
end
```

In this code, under `-gpu=mem:unified` (or `managed`) a dynamic-extent automatic array must be allocated in CUDA unified/managed memory so the host pointer is device-shared. Fixed-size locals (`ar_tmp(128)`) stay on the stack. Leaving VLAs on the ordinary host heap breaks unified-memory OpenACC device access.

Fix: under those GPU mem features, tag non-dummy explicit-shape automatic arrays with a non-constant bound as Unified/Managed in semantics so lowering uses `cuf.alloc`/`cuf.free`.


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


2 Files Affected:

- (modified) flang/lib/Semantics/resolve-names.cpp (+31) 
- (added) flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90 (+58) 


``````````diff
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index 83d92f253e626..a3514ccf72e5c 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -10648,6 +10648,37 @@ void ResolveNamesVisitor::FinishSpecificationPart(
             context().languageFeatures().IsEnabled(
                 common::LanguageFeature::CudaPinned))
           object->set_cudaDataAttr(common::CUDADataAttr::Pinned);
+      } else if (!object->cudaDataAttr() && !IsDummy(symbol) &&
+          !IsAllocatable(symbol) && !IsPointer(symbol) && !IsSaved(symbol) &&
+          !IsCUDADeviceContext(&symbol.owner()) &&
+          object->shape().IsExplicitShape()) {
+        // Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic
+        // arrays in CUDA unified/managed memory (fixed-size automatic arrays
+        // stay on the stack). Tag those locals so lowering uses
+        // cuf.alloc/cuf.free. Unlike the allocatable managed tagging above,
+        // this does not require -fcuda: OpenACC + -gpu=mem:unified relies on
+        // it, and cuf.alloc does not go through the CUDA Fortran
+        // managed-descriptor pipeline that motivated the -fcuda gate.
+        std::optional<common::CUDADataAttr> attr;
+        if (context().languageFeatures().IsEnabled(
+                common::LanguageFeature::CudaUnified))
+          attr = common::CUDADataAttr::Unified;
+        else if (context().languageFeatures().IsEnabled(
+                     common::LanguageFeature::CudaManaged))
+          attr = common::CUDADataAttr::Managed;
+        if (attr) {
+          auto boundIsNonConstant{[](const Bound &b) {
+            return !b.isExplicit() || !b.GetExplicit() ||
+                !evaluate::IsConstantExpr(*b.GetExplicit());
+          }};
+          for (const ShapeSpec &ss : object->shape()) {
+            if (boundIsNonConstant(ss.lbound()) ||
+                boundIsNonConstant(ss.ubound())) {
+              object->set_cudaDataAttr(*attr);
+              break;
+            }
+          }
+        }
       }
     }
   }
diff --git a/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90 b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
new file mode 100644
index 0000000000000..07638b0ca8dfa
--- /dev/null
+++ b/flang/test/Lower/CUDA/cuda-gpu-unified-automatic-array.f90
@@ -0,0 +1,58 @@
+! RUN: bbc -emit-hlfir -gpu=unified %s -o - | FileCheck %s
+! RUN: bbc -emit-hlfir -gpu=managed %s -o - | FileCheck %s --check-prefix=MANAGED
+
+! Under -gpu=mem:unified|managed, allocate adjustable / VLA automatic arrays
+! in CUDA unified/managed memory. Fixed-size automatic arrays stay on the
+! stack so host pointers remain shared under unified memory (e.g. OpenACC).
+
+module m_adj
+  integer :: nx = 32
+end module
+
+! CHECK-LABEL: func.func @_QPvla(
+! CHECK: %[[ALLOC:.*]] = cuf.alloc !fir.array<?xf32>, %{{.*}} : index {bindc_name = "a", data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
+! CHECK: %[[DECL:.*]]:2 = hlfir.declare %[[ALLOC]](%{{.*}}) {data_attr = #cuf.cuda<unified>, uniq_name = "_QFvlaEa"}
+! CHECK: cuf.free %[[DECL]]#1 : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<unified>}
+! MANAGED-LABEL: func.func @_QPvla(
+! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
+! MANAGED: cuf.free %{{.*}} : !fir.ref<!fir.array<?xf32>> {data_attr = #cuf.cuda<managed>}
+subroutine vla(n)
+  integer :: n
+  real :: a(n)
+  a(1) = 1.0
+end subroutine
+
+! CHECK-LABEL: func.func @_QPadjustable(
+! CHECK: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<unified>
+! CHECK: cuf.free %{{.*}} {data_attr = #cuf.cuda<unified>}
+! MANAGED-LABEL: func.func @_QPadjustable(
+! MANAGED: cuf.alloc !fir.array<?xf32>, %{{.*}} : index {{{.*}}data_attr = #cuf.cuda<managed>
+subroutine adjustable
+  use m_adj
+  real :: a(0:(nx+1)/2)
+  a(0) = 0.0
+end subroutine
+
+! Fixed-size automatic arrays must remain ordinary stack allocations.
+! CHECK-LABEL: func.func @_QPfixed(
+! CHECK-NOT: cuf.alloc
+! CHECK: fir.alloca !fir.array<128xf32>
+! CHECK-NOT: cuf.free
+! MANAGED-LABEL: func.func @_QPfixed(
+! MANAGED-NOT: cuf.alloc
+! MANAGED: fir.alloca !fir.array<128xf32>
+subroutine fixed
+  real :: a(128)
+  a(1) = 1.0
+end subroutine
+
+! Dummy adjustable arrays are caller-allocated; do not retag them.
+! CHECK-LABEL: func.func @_QPdummy_adj(
+! CHECK-NOT: cuf.alloc
+! MANAGED-LABEL: func.func @_QPdummy_adj(
+! MANAGED-NOT: cuf.alloc
+subroutine dummy_adj(a, n)
+  integer :: n
+  real :: a(n)
+  a(1) = 1.0
+end subroutine

``````````

</details>


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