[flang-commits] [flang] [llvm] [Flang][OpenMP] Clear close on descriptor members for box parents in USM (PR #194287)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: Kareem Ergawy (ergawy)

<details>
<summary>Changes</summary>

Extend the MapInfoFinalization walk introduced in #<!-- -->185330 so parent/member close consistency is enforced whenever unified_shared_memory is in effect, not only when the parent map's variable is a fir.RecordType. Allocatable (box) roots expand to member maps the same way as derived-type instances; getDescriptorMapType may add OMP_MAP_CLOSE to implicit descriptor members while the parent map does not set close, which led to bad device behavior under -fopenmp-force-usm with multiple mapped allocatables.

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


3 Files Affected:

- (modified) flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp (+6-12) 
- (modified) flang/test/Transforms/omp-map-info-finalization-usm.fir (+12-12) 
- (added) offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90 (+49) 


``````````diff
diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index bc0f96478ddf4..741d3174c29ee 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1186,21 +1186,15 @@ class MapInfoFinalizationPass
       });
 
       func->walk([&](mlir::omp::MapInfoOp op) {
-        // If a record type is not mapped with the `close` modifier while some
-        // of its members are (e.g. descriptor maps), then in USM mode, the
-        // memory for the record will be allocated in unified memory while the
-        // the members might be allocated in device memory. This creates an
-        // inconsistent map for the record type where some of its members are
-        // allocated in different address spaces.
-        //
-        // This fixes this issue by taking a conservative approach and removing
-        // the `close` flag from members if it is not used for mapping the
-        // parent record.
+        // If a parent map is not mapped with the `close` modifier while some of
+        // its members are (e.g. implicit descriptor maps from
+        // getDescriptorMapType in USM), those members must not keep `close` —
+        // otherwise the runtime can treat unified and device placement
+        // inconsistently.
         if (op.getMembers().empty())
           return;
 
-        mlir::Type varTy = fir::unwrapRefType(op.getVarPtr().getType());
-        if (!mlir::isa<fir::RecordType>(varTy))
+        if (!moduleRequiresUSM(op->getParentOfType<mlir::ModuleOp>()))
           return;
 
         auto mapFlag = op.getMapType();
diff --git a/flang/test/Transforms/omp-map-info-finalization-usm.fir b/flang/test/Transforms/omp-map-info-finalization-usm.fir
index 5f5a0d7213719..24f08474ed1d5 100644
--- a/flang/test/Transforms/omp-map-info-finalization-usm.fir
+++ b/flang/test/Transforms/omp-map-info-finalization-usm.fir
@@ -1,24 +1,24 @@
 // RUN: fir-opt --split-input-file --omp-map-info-finalization %s | FileCheck %s
 
 // Test that the 'close' map flag is cleared from member maps if the parent map
-// (derived type) does not have the 'close' flag. This typically happens in
-// Unified Shared Memory (USM) mode where the parent is in USM (no close) but
-// members (like descriptors) might have been initially tagged with close.
-
+// does not have the 'close' flag. This typically happens in USM mode.
+// In this test, the parent is a fir.array of derived type.
 module attributes {omp.requires = #omp<clause_requires unified_shared_memory>} {
-  func.func @test_usm_close_flag_cleanup(%arg0: !fir.ref<!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>) {
-    %map = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>, !fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>) map_clauses(to) capture(ByRef) -> !fir.ref<!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>> {name = "parent"}
+  func.func @test_usm_close_flag_cleanup_array(%arg0: !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>>) {
+    // The implicit descriptor map for the member, with 'close'
+    %member = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>>, !fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>) map_clauses(always, close, descriptor, to) capture(ByRef) -> !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>> {name = "parent.a.implicit_map"}
+
+    // The parent map, which is a fir.array, without 'close'
+    %map = omp.map.info var_ptr(%arg0 : !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>>, !fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>) map_clauses(to) capture(ByRef) members(%member : [0] : !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>>) -> !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>> {name = "parent"}
 
-    omp.target map_entries(%map -> %arg1 : !fir.ref<!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>) {
-      // Simulate usage to trigger implicit map addition
-      %1 = hlfir.designate %arg1{"a"} : (!fir.ref<!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    omp.target map_entries(%map -> %arg1 : !fir.ref<!fir.array<10x!fir.type<t{a:!fir.box<!fir.heap<!fir.array<?xf32>>>}>>>) {
       omp.terminator
     }
     return
   }
 }
 
-// CHECK-LABEL: func.func @test_usm_close_flag_cleanup
-// CHECK: %[[MEMBER:.*]] = omp.map.info {{.*}} map_clauses(always, to) {{.*}} {name = "parent.a.implicit_map"}
-// CHECK: %[[PARENT:.*]] = omp.map.info {{.*}} map_clauses(to) {{.*}} members(%[[MEMBER]], {{.*}}) {{.*}} {name = "parent", {{.*}}}
+// CHECK-LABEL: func.func @test_usm_close_flag_cleanup_array
+// CHECK: %[[MEMBER:.*]] = omp.map.info {{.*}} map_clauses(always, to) {{.*}} {name = "parent{{.*}}implicit_map"}
+// CHECK: %[[PARENT:.*]] = omp.map.info {{.*}} map_clauses(to) {{.*}} members(%[[MEMBER]]{{.*}}) {{.*}} {name = "parent"}
 // CHECK-NOT: close
diff --git a/offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90 b/offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90
new file mode 100644
index 0000000000000..6bacc1b4a761b
--- /dev/null
+++ b/offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90
@@ -0,0 +1,49 @@
+! Test for PR fixing close flag on descriptor members for box parents in USM
+! REQUIRES: flang, amdgpu
+! RUN: %libomptarget-compile-fortran-generic -fopenmp-force-usm
+! RUN: env LIBOMPTARGET_INFO=16 HSA_XNACK=1 %libomptarget-run-generic 2>&1 | %fcheck-generic
+
+module m
+    implicit none
+    integer :: ng
+    type :: gt
+        integer :: k
+    end type
+    type(gt), allocatable :: g(:)
+    !$omp declare target(ng, g)
+    type :: f
+        real, allocatable :: a(:)
+    end type
+end module m
+
+program r
+    use m
+    implicit none
+    integer :: i
+    type(f), target :: u(2)
+    integer :: ig
+    real, contiguous, pointer :: p(:)
+
+    ng = 1
+    allocate(g(1))
+    g(1)%k = 1
+
+    do i = 1, 2
+        allocate(u(i)%a(1), source=0.0)
+    end do
+    u(1)%a(1) = 1.0
+    u(2)%a(1) = -1.0
+
+    !$omp target enter data map(to: g, ng, u(1)%a, u(2)%a)
+
+    !$omp target teams distribute private(ig, p)
+    do ig = 1, ng
+        p(1:1) => u(2)%a(1:1)
+        p(1) = 3.14
+    end do
+    !$omp end target teams distribute
+
+    ! CHECK: PluginInterface device {{[0-9]+}} info: Launching kernel
+    ! CHECK: Result: 3.14
+    print *, "Result: ", u(2)%a(1)
+end program r

``````````

</details>


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