[flang-commits] [flang] [llvm] [Flang][OpenMP] Fix close map flag propagation for derived types in USM (PR #185330)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: Kareem Ergawy (ergawy)

<details>
<summary>Changes</summary>

This fixes a bug in USM mode where the `close` map type modifer was attached to some `map.info.op`'s corresponding to user-defined type members while the parent type instance itself is not marked as `close`.

This fix ensures that if a parent record type map does not have the 'close' flag, it is cleared from its members as well, maintaining consistency.

Gemini was used to create tests. AI generated test code was reviewed line-by-line by me. Which were derived from a reproducer I was working with to debug the issue.

Assisted-by: Gemini <gemini@<!-- -->google.com>

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


3 Files Affected:

- (modified) flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp (+35) 
- (added) flang/test/Transforms/omp-map-info-finalization-usm.fir (+24) 
- (added) offload/test/offloading/fortran/usm_derived_type_allocatable_member.f90 (+35) 


``````````diff
diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index a60960e739d24..bc0f96478ddf4 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -1185,6 +1185,41 @@ 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 (op.getMembers().empty())
+          return;
+
+        mlir::Type varTy = fir::unwrapRefType(op.getVarPtr().getType());
+        if (!mlir::isa<fir::RecordType>(varTy))
+          return;
+
+        auto mapFlag = op.getMapType();
+        bool hasClose = (mapFlag & mlir::omp::ClauseMapFlags::close) ==
+                        mlir::omp::ClauseMapFlags::close;
+
+        if (hasClose)
+          return;
+
+        for (auto member : op.getMembers()) {
+          if (auto memberOp = llvm::dyn_cast_if_present<mlir::omp::MapInfoOp>(
+                  member.getDefiningOp())) {
+            auto memberMapFlag =
+                memberOp.getMapType() & ~mlir::omp::ClauseMapFlags::close;
+            memberOp.setMapType(memberMapFlag);
+          }
+        }
+      });
+
       // Now that we've expanded all of our boxes into a descriptor and base
       // address map where necessary, we check if the map owner is an
       // enter/exit/target data directive, and if they are we drop the initial
diff --git a/flang/test/Transforms/omp-map-info-finalization-usm.fir b/flang/test/Transforms/omp-map-info-finalization-usm.fir
new file mode 100644
index 0000000000000..5f5a0d7213719
--- /dev/null
+++ b/flang/test/Transforms/omp-map-info-finalization-usm.fir
@@ -0,0 +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.
+
+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"}
+
+    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.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-NOT: close
diff --git a/offload/test/offloading/fortran/usm_derived_type_allocatable_member.f90 b/offload/test/offloading/fortran/usm_derived_type_allocatable_member.f90
new file mode 100644
index 0000000000000..1653fc38f3a41
--- /dev/null
+++ b/offload/test/offloading/fortran/usm_derived_type_allocatable_member.f90
@@ -0,0 +1,35 @@
+! Offloading test checking USM with derived types containing allocatable members.
+! This ensures that the 'close' map flag is not incorrectly applied to members
+! when the parent is in USM, preventing runtime crashes.
+! 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 globalmod
+   IMPLICIT NONE
+   TYPE :: GRID_type
+      REAL(KIND=8),ALLOCATABLE,DIMENSION(:) :: DZ
+      REAL(KIND=8),ALLOCATABLE,DIMENSION(:) :: RDZ
+   END TYPE GRID_type
+   TYPE(GRID_type) :: GRID
+END MODULE globalmod
+
+PROGRAM reproducer
+   USE globalmod
+   IMPLICIT NONE
+
+   ALLOCATE(GRID%DZ(10))
+   ALLOCATE(GRID%RDZ(10))
+   GRID%DZ = 1.0
+   GRID%RDZ = 2.0
+
+   !$OMP TARGET
+   GRID%DZ(1) = GRID%DZ(1) + GRID%RDZ(1)
+   !$OMP END TARGET
+
+   PRINT *, GRID%DZ(1)
+END PROGRAM reproducer
+
+! CHECK: PluginInterface device {{[0-9]+}} info: Launching kernel
+! CHECK: 3.

``````````

</details>


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