diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp index 741d3174c29ee..bc0f96478ddf4 100644 --- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp +++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp @@ -1186,15 +1186,21 @@ class MapInfoFinalizationPass }); func->walk([&](mlir::omp::MapInfoOp op) { - // 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 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; - if (!moduleRequiresUSM(op->getParentOfType())) + mlir::Type varTy = fir::unwrapRefType(op.getVarPtr().getType()); + if (!mlir::isa(varTy)) 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 24f08474ed1d5..5f5a0d7213719 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 -// 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} { - func.func @test_usm_close_flag_cleanup_array(%arg0: !fir.ref>>}>>>) { - // The implicit descriptor map for the member, with 'close' - %member = omp.map.info var_ptr(%arg0 : !fir.ref>>}>>>, !fir.array<10x!fir.type>>}>>) map_clauses(always, close, descriptor, to) capture(ByRef) -> !fir.ref>>}>>> {name = "parent.a.implicit_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. - // The parent map, which is a fir.array, without 'close' - %map = omp.map.info var_ptr(%arg0 : !fir.ref>>}>>>, !fir.array<10x!fir.type>>}>>) map_clauses(to) capture(ByRef) members(%member : [0] : !fir.ref>>}>>>) -> !fir.ref>>}>>> {name = "parent"} +module attributes {omp.requires = #omp} { + func.func @test_usm_close_flag_cleanup(%arg0: !fir.ref>>}>>) { + %map = omp.map.info var_ptr(%arg0 : !fir.ref>>}>>, !fir.type>>}>) map_clauses(to) capture(ByRef) -> !fir.ref>>}>> {name = "parent"} - omp.target map_entries(%map -> %arg1 : !fir.ref>>}>>>) { + omp.target map_entries(%map -> %arg1 : !fir.ref>>}>>) { + // Simulate usage to trigger implicit map addition + %1 = hlfir.designate %arg1{"a"} : (!fir.ref>>}>>) -> !fir.ref>>> omp.terminator } return } } -// 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-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-box-parent-descriptor-close.f90 b/offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90 deleted file mode 100644 index 6bacc1b4a761b..0000000000000 --- a/offload/test/offloading/fortran/usm-box-parent-descriptor-close.f90 +++ /dev/null @@ -1,49 +0,0 @@ -! 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