[flang-commits] [flang] [mlir] [Flang][OpenMP] Privatize descriptors for assumed shape array maps for performance increase (PR #212336)

via flang-commits flang-commits at lists.llvm.org
Mon Jul 27 12:50:05 PDT 2026


https://github.com/agozillon created https://github.com/llvm/llvm-project/pull/212336

This PR aims to decrease the performance overhead of descriptor mapping by privatizing the descriptors, thus having them part of the initial kernel payload as opposed to a separate more costly H2D transfer. We do so by modifying the map types for the descriptor, hooking into the runtimes existing privatization of attach pointers. Some minor tweaks in the lowering to LLVM-IR are also required to specialize around this attach map privatization case. The intent is to expand this privatization to more comprehensively be the default case for descriptors, so hopefully we'll be able to assimilate the edge case better with the member mapping in the future.

This currently only applies to assumed shape array arguments while we test the cost effectiveness and possible downsides.

Co-author: Akash Banerjee <Akash.Banerjee at amd.com>

>From 7259d4106aa5e946d380ae3424086907213974d9 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Tue, 14 Jul 2026 22:08:35 +0200
Subject: [PATCH] [Flang][OpenMP] Privatize descriptors for assumed-size maps
 for performance increase (#3342)

This PR aims to decrease the performance overhead of descriptor mapping
by privatizing the descriptors, thus having them part of the initial
kernel payload as opposd to a seperate more costly H2D transfer.

This currently only applies to assumed size variables/arguments while we
test the cost effectiveness and possible downsides.

Co-author: Akash Banerjee <Akash.Banerjee at amd.com>
---
 .../Optimizer/OpenMP/MapInfoFinalization.cpp  | 76 ++++++++++++++-----
 flang/test/Lower/OpenMP/array-bounds.f90      |  2 +-
 .../Lower/OpenMP/map-descriptor-deferral.f90  |  4 +-
 .../OpenMP/map-descriptor-privatization.f90   | 19 +++++
 .../Lower/OpenMP/optional-argument-map-3.f90  |  4 +-
 .../mlir/Dialect/OpenMP/OpenMPEnums.td        |  4 +-
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  |  5 ++
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 76 ++++++++++++++-----
 .../omptarget-map-pointer-privatization.mlir  | 39 ++++++++++
 9 files changed, 184 insertions(+), 45 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/map-descriptor-privatization.f90
 create mode 100644 mlir/test/Target/LLVMIR/omptarget-map-pointer-privatization.mlir

diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index 992ca8b94c5cb..fba190b5a69bc 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -340,7 +340,8 @@ class MapInfoFinalizationPass
   /// base address can be utilised.
   mlir::Value getDescriptorFromBoxMap(mlir::omp::MapInfoOp boxMap,
                                       fir::FirOpBuilder &builder,
-                                      bool &canDescBeDeferred) {
+                                      bool &canDescBeDeferred,
+                                      bool &canOptimizeDescViaPrivatization) {
     mlir::Value descriptor = boxMap.getVarPtr();
     if (!fir::isTypeWithDescriptor(boxMap.getVarPtrType()))
       if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
@@ -349,8 +350,13 @@ class MapInfoFinalizationPass
 
     canDescBeDeferred = canDeferDescriptorMapping(descriptor);
 
-    if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
-        !fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
+    // A restricted subset of canDeferDescriptorMapping for the moment
+    // can extend to encompass more as we find acceptable cases.
+    canOptimizeDescViaPrivatization = isDummyArgument(descriptor) &&
+                                      fir::isAssumedShape(descriptor.getType());
+
+    if ((!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
+         !fir::factory::isOptionalArgument(descriptor.getDefiningOp())))
       return descriptor;
 
     mlir::Value &alloca = localBoxAllocas[descriptor.getDefiningOp()];
@@ -752,10 +758,20 @@ class MapInfoFinalizationPass
   /// issues.
   mlir::omp::ClauseMapFlags
   getDescriptorMapType(mlir::omp::ClauseMapFlags mapTypeFlag,
-                       mlir::Operation *target) {
+                       mlir::Operation *target, bool privatizeDescriptor) {
     using MapFlags = mlir::omp::ClauseMapFlags;
     MapFlags flags = MapFlags::none;
 
+    // Special runtime case for descriptor privatization requires the
+    // following map types in synergy:
+    //
+    //  PRIVATE | ATTACH | TARGET_PARAM
+    //
+    if (privatizeDescriptor) {
+      return MapFlags::priv | MapFlags::attach | MapFlags::target_param |
+             (mapTypeFlag & MapFlags::implicit);
+    }
+
     if (llvm::isa_and_nonnull<mlir::omp::TargetExitDataOp,
                               mlir::omp::TargetUpdateOp>(target)) {
       return mapTypeFlag;
@@ -817,6 +833,18 @@ class MapInfoFinalizationPass
     return false;
   }
 
+  /// Gets the underlying type of a pointer type, effectively unwrapping
+  /// fir.ref, and fir.array to get the underlying scalar type.
+  mlir::Type getUnderlyingVarType(mlir::Type baseAddrType) {
+    baseAddrType =
+        llvm::cast<mlir::omp::PointerLikeType>(fir::unwrapRefType(baseAddrType))
+            .getElementType();
+    if (auto seqType = llvm::dyn_cast<fir::SequenceType>(baseAddrType))
+      if (seqType.hasDynamicExtents())
+        baseAddrType = seqType.getEleTy();
+    return baseAddrType;
+  }
+
   /// This function generates an attach map, which is an type of OpenMP map that
   /// binds a pointer to its data. In the case of Fortran, this binding is
   /// primarily for binding the pointer inside of descriptors to the underlying
@@ -838,12 +866,8 @@ class MapInfoFinalizationPass
             ? reuseBaseAddr
             : fir::BoxOffsetOp::create(builder, descMapOp->getLoc(), descriptor,
                                        fir::BoxFieldAttr::base_addr);
-    mlir::Type underlyingVarType = llvm::cast<mlir::omp::PointerLikeType>(
-                                       fir::unwrapRefType(baseAddr.getType()))
-                                       .getElementType();
-    if (auto seqType = llvm::dyn_cast<fir::SequenceType>(underlyingVarType))
-      if (seqType.hasDynamicExtents())
-        underlyingVarType = seqType.getEleTy();
+
+    mlir::Type underlyingVarType = getUnderlyingVarType(baseAddr.getType());
 
     auto implicitAttachMap = mlir::omp::MapInfoOp::create(
         builder, descMapOp->getLoc(), descMapOp.getResult().getType(),
@@ -1068,7 +1092,8 @@ class MapInfoFinalizationPass
       mlir::Operation *target, mlir::Value descriptor,
       llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers,
       bool isAttachNever, bool isAttachAlways, bool isHasDeviceAddrFlag,
-      bool descCanBeDeferred, mlir::FlatSymbolRefAttr mapperId) {
+      bool descCanBeDeferred, bool canOptimizeDescViaPrivatization,
+      mlir::FlatSymbolRefAttr mapperId) {
     bool isRefPtrPtee =
         bitEnumContainsAll(op.getMapType(),
                            mlir::omp::ClauseMapFlags::ref_ptr) &&
@@ -1095,19 +1120,30 @@ class MapInfoFinalizationPass
                               newMembersAttr, newMembers, memberIndices);
     }
 
+    bool optDescMap = canOptimizeDescViaPrivatization &&
+                      llvm::isa<mlir::omp::TargetOp>(target);
+
     // If we have been provided RefPtrPtee, utilise the user specified map
     // types, otherwise, use the default descriptor map types.
     auto mapType = isRefPtrPtee ? op.getMapType()
-                                : getDescriptorMapType(op.getMapType(), target);
+                                : getDescriptorMapType(op.getMapType(), target,
+                                                       optDescMap);
 
     mapType = removeAttachModifiers(mapType);
 
+    mlir::Type underlyingVarType = mlir::Type{};
+    if (optDescMap)
+      underlyingVarType = getUnderlyingVarType(baseAddr.getType());
+
     auto newMapInfoOp = mlir::omp::MapInfoOp::create(
         builder, op->getLoc(), op.getResult().getType(), descriptor,
         mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
         builder.getAttr<mlir::omp::ClauseMapFlagsAttr>(mapType),
-        op.getMapCaptureTypeAttr(), /*varPtrPtr=*/mlir::Value{},
-        /*varPtrPtTyper=*/mlir::TypeAttr{}, newMembers, newMembersAttr,
+        op.getMapCaptureTypeAttr(),
+        optDescMap ? baseAddr.getVarPtrPtr() : mlir::Value{},
+        underlyingVarType ? mlir::TypeAttr::get(underlyingVarType)
+                          : mlir::TypeAttr{},
+        newMembers, newMembersAttr,
         /*bounds=*/mlir::SmallVector<mlir::Value>{},
         /*mapperId*/ mlir::FlatSymbolRefAttr(), op.getNameAttr(),
         /*partial_map=*/builder.getBoolAttr(false));
@@ -1143,6 +1179,7 @@ class MapInfoFinalizationPass
   void genDescriptorMaps(mlir::omp::MapInfoOp op, fir::FirOpBuilder &builder,
                          mlir::Operation *target) {
     bool descCanBeDeferred = false;
+    bool canOptimizeDescViaPrivatization = false;
     llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
     getMemberUserList(op, mapMemberUsers);
 
@@ -1162,8 +1199,8 @@ class MapInfoFinalizationPass
                      !bitEnumContainsAll(op.getMapType(),
                                          mlir::omp::ClauseMapFlags::ref_ptr);
 
-    mlir::Value descriptor =
-        getDescriptorFromBoxMap(op, builder, descCanBeDeferred);
+    mlir::Value descriptor = getDescriptorFromBoxMap(
+        op, builder, descCanBeDeferred, canOptimizeDescViaPrivatization);
     mlir::FlatSymbolRefAttr mapperId = op.getMapperIdAttr();
 
     // If we're a derived type descriptor, that's been flagged as ref_ptr,
@@ -1189,9 +1226,10 @@ class MapInfoFinalizationPass
       genRefPteeMap(op, builder, target, descriptor, mapMemberUsers,
                     isAttachNever, isAttachAlways, mapperId);
     } else {
-      genRefPtrPteeOrDefaultMap(
-          op, builder, target, descriptor, mapMemberUsers, isAttachNever,
-          isAttachAlways, isHasDeviceAddrFlag, descCanBeDeferred, mapperId);
+      genRefPtrPteeOrDefaultMap(op, builder, target, descriptor, mapMemberUsers,
+                                isAttachNever, isAttachAlways,
+                                isHasDeviceAddrFlag, descCanBeDeferred,
+                                canOptimizeDescViaPrivatization, mapperId);
     }
   }
 
diff --git a/flang/test/Lower/OpenMP/array-bounds.f90 b/flang/test/Lower/OpenMP/array-bounds.f90
index a9f2ad448c632..382e559c673ee 100644
--- a/flang/test/Lower/OpenMP/array-bounds.f90
+++ b/flang/test/Lower/OpenMP/array-bounds.f90
@@ -52,7 +52,7 @@ module assumed_array_routines
 !HOST: %[[BOUNDS:.*]] = omp.map.bounds   lower_bound(%[[C3]] : index) upper_bound(%[[C4]] : index) extent(%[[DIMS1]]#1 : index) stride(%[[DIMS0]]#2 : index) start_idx(%[[C0]] : index) {stride_in_bytes = true}
 !HOST: %[[VAR_PTR_PTR:.*]] = fir.box_offset %{{.*}} base_addr : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
 !HOST: %[[MAP_INFO_MEMBER:.*]] = omp.map.info var_ptr(%[[INTERMEDIATE_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(tofrom) capture(ByRef) var_ptr_ptr(%[[VAR_PTR_PTR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32)   bounds(%[[BOUNDS]]) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>> {name = ""}
-!HOST: %[[MAP:.*]] = omp.map.info var_ptr(%[[INTERMEDIATE_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, to) capture(ByRef) members(%[[MAP_INFO_MEMBER]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "arr_read_write(2:5)"}
+!HOST: %[[MAP:.*]] = omp.map.info var_ptr(%[[INTERMEDIATE_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(target_param, private, attach) capture(ByRef) var_ptr_ptr(%[[VAR_PTR_PTR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) members(%[[MAP_INFO_MEMBER]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "arr_read_write(2:5)"}
 !HOST: %[[ATTACH_MAP:.*]] = omp.map.info var_ptr(%[[INTERMEDIATE_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(attach, ref_ptr, ref_ptee) capture(ByRef) var_ptr_ptr(%[[VAR_PTR_PTR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%[[BOUNDS]]) -> !fir.ref<!fir.array<?xi32>> {name = "arr_read_write(2:5)"}
 !HOST: omp.target kernel_type(generic) map_entries(%[[MAP]] -> %{{.*}}, {{.*}} -> {{.*}}, %[[ATTACH_MAP]] -> {{.*}}, %[[MAP_INFO_MEMBER]] -> %{{.*}} : !fir.ref<!fir.array<?xi32>>, !fir.ref<i32>, !fir.ref<!fir.array<?xi32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) {
     subroutine assumed_shape_array(arr_read_write)
diff --git a/flang/test/Lower/OpenMP/map-descriptor-deferral.f90 b/flang/test/Lower/OpenMP/map-descriptor-deferral.f90
index fee20612d8bc3..35398f4886dad 100644
--- a/flang/test/Lower/OpenMP/map-descriptor-deferral.f90
+++ b/flang/test/Lower/OpenMP/map-descriptor-deferral.f90
@@ -23,7 +23,7 @@ subroutine assume_map_target_enter_exit(assumed_arr)
 !CHECK:    omp.target_enter_data map_entries(%[[MAP_ADDR]] : !fir.ref<!fir.array<?xi32>>)
 !CHECK:    %[[BOX_ADDR:.*]] = fir.box_offset %{{.*}} base_addr : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
 !CHECK:    %[[MAP_ADDR:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(implicit, tofrom) capture(ByRef) var_ptr_ptr(%[[BOX_ADDR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>> {name = ""}
-!CHECK:    %[[MAP_BOX:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, implicit, to) capture(ByRef) members(%{{.*}} : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
+!CHECK:    %[[MAP_BOX:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(implicit, target_param, private, attach) capture(ByRef) var_ptr_ptr(%{{.*}} : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) members(%{{.*}} : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
 !CHECK:    %[[MAP_BOX_ATTACH:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(attach, ref_ptr, ref_ptee) capture(ByRef) var_ptr_ptr(%{{.*}} : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
 !CHECK:    omp.target kernel_type(generic) map_entries(%[[MAP_BOX]] -> %{{.*}}, %[[MAP_BOX_ATTACH]] -> %{{.*}}, %[[MAP_ADDR]] -> %{{.*}} : !fir.ref<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) {
 !CHECK:    %[[BOX_ADDR:.*]] = fir.box_offset %{{.*}} base_addr : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
@@ -100,6 +100,6 @@ subroutine assume_map_target_data(assumed_arr)
 !CHECK:    omp.target_data map_entries(%[[MAP_BOX]], %[[ATTACH]], %[[MAP_ADDR]] : !fir.ref<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) {
 !CHECK:    %[[BOX_ADDR:.*]] = fir.box_offset %{{.*}} base_addr : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
 !CHECK:    %[[MAP_ADDR:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(implicit, tofrom) capture(ByRef) var_ptr_ptr(%[[BOX_ADDR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>> {name = ""}
-!CHECK:    %[[MAP_BOX:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(always, implicit, to) capture(ByRef) members(%[[MAP_ADDR]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
+!CHECK:    %[[MAP_BOX:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(implicit, target_param, private, attach) capture(ByRef) var_ptr_ptr(%{{.*}} : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) members(%[[MAP_ADDR]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
 !CHECK:    %[[ATTACH:.*]] = omp.map.info var_ptr(%{{.*}} : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(attach, ref_ptr, ref_ptee) capture(ByRef) var_ptr_ptr(%{{.*}} : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.ref<!fir.array<?xi32>> {name = "assumed_arr"}
 !CHECK:    omp.target kernel_type(generic) map_entries(%[[MAP_BOX]] -> %{{.*}}, %[[ATTACH]] -> %{{.*}}, %[[MAP_ADDR]] -> %{{.*}} : !fir.ref<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) {
diff --git a/flang/test/Lower/OpenMP/map-descriptor-privatization.f90 b/flang/test/Lower/OpenMP/map-descriptor-privatization.f90
new file mode 100644
index 0000000000000..ffa36d19bc263
--- /dev/null
+++ b/flang/test/Lower/OpenMP/map-descriptor-privatization.f90
@@ -0,0 +1,19 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+! This test checks the descriptor privatization for assumed arrays verifying the
+! maps have the appropriate map types applied to undergo attach map privatization.
+
+subroutine assumed_shape_array_priv(arr_read_write)
+    integer, intent(inout) :: arr_read_write(:)
+    !$omp target map(tofrom: arr_read_write)
+        arr_read_write(1) = 10
+    !$omp end target
+end subroutine
+
+!CHECK-LABEL:   func.func @_QPassumed_shape_array_priv(
+!CHECK:    %[[DESC_ALLOCA:.*]] = fir.alloca !fir.box<!fir.array<?xi32>>
+!CHECK:    %[[BOX_ADDR:.*]] = fir.box_offset %[[DESC_ALLOCA]] base_addr : (!fir.ref<!fir.box<!fir.array<?xi32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+!CHECK:    %[[MAP_MEMBER:.*]] = omp.map.info var_ptr(%[[DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(tofrom) capture(ByRef) var_ptr_ptr(%[[BOX_ADDR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>> {name = ""}
+!CHECK:    %[[MAP_PARENT:.*]] = omp.map.info var_ptr(%[[DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(target_param, private, attach) capture(ByRef) var_ptr_ptr(%[[BOX_ADDR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) members(%[[MAP_MEMBER]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) -> !fir.ref<!fir.array<?xi32>> {name = "arr_read_write"}
+!CHECK:    %[[MAP_ATTACH:.*]] = omp.map.info var_ptr(%[[DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xi32>>>, !fir.box<!fir.array<?xi32>>) map_clauses(attach, ref_ptr, ref_ptee) capture(ByRef) var_ptr_ptr(%[[BOX_ADDR]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>, i32) bounds(%{{.*}}) -> !fir.ref<!fir.array<?xi32>> {name = "arr_read_write"}
+!CHECK:    omp.target kernel_type(generic) map_entries(%[[MAP_PARENT]] -> %{{.*}}, %[[MAP_ATTACH]] -> %{{.*}}, %[[MAP_MEMBER]] -> %{{.*}} : !fir.ref<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>) {
diff --git a/flang/test/Lower/OpenMP/optional-argument-map-3.f90 b/flang/test/Lower/OpenMP/optional-argument-map-3.f90
index 2a0c78d50af96..a7b7b06f21bfd 100644
--- a/flang/test/Lower/OpenMP/optional-argument-map-3.f90
+++ b/flang/test/Lower/OpenMP/optional-argument-map-3.f90
@@ -33,7 +33,7 @@ end subroutine foo
 ! CHECK:             }
 ! CHECK:             %[[VAL_3:.*]] = fir.box_offset %[[VAL_0]] base_addr : (!fir.ref<!fir.box<!fir.array<?xf32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>
 ! CHECK:             %[[VAL_4:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(implicit, tofrom) capture(ByRef) var_ptr_ptr(%[[VAL_3]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>, f32) bounds(%{{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>> {name = ""}
-! CHECK:             %[[VAL_5:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(always, implicit, to) capture(ByRef) members(%[[VAL_4]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>) -> !fir.ref<!fir.array<?xf32>> {name = "dt"}
+! CHECK:             %[[VAL_5:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(implicit, target_param, private, attach) capture(ByRef) var_ptr_ptr(%[[VAL_3]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>, f32) members(%[[VAL_4]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>) -> !fir.ref<!fir.array<?xf32>> {name = "dt"}
 ! CHECK:             omp.target kernel_type(spmd) host_eval({{.*}}) map_entries({{.*}}%[[VAL_5]] -> {{.*}}, %[[VAL_4]] -> {{.*}} : {{.*}}) {
 ! CHECK:           } else {
 ! CHECK:             %[[VAL_6:.*]] = fir.is_present %[[VAL_1]]#1 : (!fir.box<!fir.array<?xf32>>) -> i1
@@ -42,5 +42,5 @@ end subroutine foo
 ! CHECK:             }
 ! CHECK:             %[[VAL_7:.*]] = fir.box_offset %[[VAL_0]] base_addr : (!fir.ref<!fir.box<!fir.array<?xf32>>>) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>
 ! CHECK:             %[[VAL_8:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(implicit, tofrom) capture(ByRef) var_ptr_ptr(%[[VAL_7]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>, f32) bounds(%{{.*}}) -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>> {name = ""}
-! CHECK:             %[[VAL_9:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(always, implicit, to) capture(ByRef) members(%[[VAL_8]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>) -> !fir.ref<!fir.array<?xf32>> {name = "dt"}
+! CHECK:             %[[VAL_9:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) map_clauses(implicit, target_param, private, attach) capture(ByRef) var_ptr_ptr(%[[VAL_7]] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>, f32) members(%[[VAL_8]] : [0] : !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>) -> !fir.ref<!fir.array<?xf32>> {name = "dt"}
 ! CHECK:             omp.target kernel_type(spmd) host_eval({{.*}}) map_entries({{.*}}, %[[VAL_9]] ->{{.*}}, %[[VAL_8]] -> {{.*}} : {{.*}}) {
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
index c642792f67019..f39f0d00ad6e1 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
@@ -148,6 +148,7 @@ def ClauseMapFlagsAttachAuto : I32BitEnumAttrCaseBit<"attach_auto", 15>;
 def ClauseMapFlagsRefPtr : I32BitEnumAttrCaseBit<"ref_ptr", 16>;
 def ClauseMapFlagsRefPtee : I32BitEnumAttrCaseBit<"ref_ptee", 17>;
 def ClauseMapFlagsIsDevicePtr : I32BitEnumAttrCaseBit<"is_device_ptr", 18>;
+def ClauseMapFlagsTargetParam : I32BitEnumAttrCaseBit<"target_param", 19>;
 
 def ClauseMapFlags : OpenMP_BitEnumAttr<
     "ClauseMapFlags",
@@ -171,7 +172,8 @@ def ClauseMapFlags : OpenMP_BitEnumAttr<
       ClauseMapFlagsAttachAuto,
       ClauseMapFlagsRefPtr,
       ClauseMapFlagsRefPtee,
-      ClauseMapFlagsIsDevicePtr
+      ClauseMapFlagsIsDevicePtr,
+      ClauseMapFlagsTargetParam
     ]>;
 
 def ClauseMapFlagsAttr : OpenMP_EnumAttr<ClauseMapFlags,
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 8605b39f17c16..bb4fa03e1d767 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2212,6 +2212,9 @@ static ParseResult parseMapClause(OpAsmParser &parser,
     if (mapTypeMod == "is_device_ptr")
       mapTypeBits |= ClauseMapFlags::is_device_ptr;
 
+    if (mapTypeMod == "target_param")
+      mapTypeBits |= ClauseMapFlags::target_param;
+
     return success();
   };
 
@@ -2243,6 +2246,8 @@ static void printMapClause(OpAsmPrinter &p, Operation *op,
     mapTypeStrs.push_back("close");
   if (mapTypeToBool(mapFlags, ClauseMapFlags::present))
     mapTypeStrs.push_back("present");
+  if (mapTypeToBool(mapFlags, ClauseMapFlags::target_param))
+    mapTypeStrs.push_back("target_param");
 
   // special handling of to/from/tofrom/delete and release/alloc, release +
   // alloc are the abscense of one of the other flags, whereas tofrom requires
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 1c50ff192c3d5..90b0df8936093 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6457,6 +6457,16 @@ static bool checkIfPointerMap(omp::MapInfoOp mapOp) {
   return false;
 }
 
+// A privatizeable attach map is a pointer/descriptor that is privatized and
+// passed directly as a kernel argument (target_param) rather than undergoing
+// the standard attach/parent mapping. These are handled specially in a couple
+// of places in map lowering.
+static bool isPrivatizeableAttachMap(omp::ClauseMapFlags mapType) {
+  return bitEnumContainsAll(mapType, omp::ClauseMapFlags::priv |
+                                         omp::ClauseMapFlags::target_param |
+                                         omp::ClauseMapFlags::attach);
+}
+
 // This function calculates the size to be offloaded for a specified type, given
 // its associated map clause (which can contain bounds information which affects
 // the total size), this size is calculated based on the underlying element type
@@ -6609,6 +6619,9 @@ convertClauseMapFlags(omp::ClauseMapFlags mlirFlags) {
   if (bitEnumContainsAll(mlirFlags, omp::ClauseMapFlags::attach))
     mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ATTACH;
 
+  if (bitEnumContainsAll(mlirFlags, omp::ClauseMapFlags::target_param))
+    mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TARGET_PARAM;
+
   if (bitEnumContainsAll(mlirFlags, omp::ClauseMapFlags::is_device_ptr)) {
     mapType |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TARGET_PARAM;
     if (!hasExplicitMap)
@@ -6650,19 +6663,25 @@ static void collectMapDataFromMapOperands(
     return false;
   };
 
+  auto checkPrivatizeableAttachPointer = [](omp::ClauseMapFlags mapType) {
+    return bitEnumContainsAll(mapType, omp::ClauseMapFlags::priv) &&
+           bitEnumContainsAll(mapType, omp::ClauseMapFlags::target_param) &&
+           bitEnumContainsAll(mapType, omp::ClauseMapFlags::attach);
+  };
+
   // Process MapOperands
   for (Value mapValue : mapVars) {
     auto mapOp = cast<omp::MapInfoOp>(mapValue.getDefiningOp());
-    bool isRefPtrOrPteeMapWithAttach =
-        checkRefPtrOrPteeMapWithAttach(mapOp.getMapType());
-    Value offloadPtr = (mapOp.getVarPtrPtr() && !isRefPtrOrPteeMapWithAttach)
+    bool isAttachStyleMap =
+        checkRefPtrOrPteeMapWithAttach(mapOp.getMapType()) ||
+        checkPrivatizeableAttachPointer(mapOp.getMapType());
+    Value offloadPtr = (mapOp.getVarPtrPtr() && !isAttachStyleMap)
                            ? mapOp.getVarPtrPtr()
                            : mapOp.getVarPtr();
     mapData.OriginalValue.push_back(moduleTranslation.lookupValue(offloadPtr));
     mapData.Pointers.push_back(
-        isRefPtrOrPteeMapWithAttach
-            ? moduleTranslation.lookupValue(mapOp.getVarPtrPtr())
-            : mapData.OriginalValue.back());
+        isAttachStyleMap ? moduleTranslation.lookupValue(mapOp.getVarPtrPtr())
+                         : mapData.OriginalValue.back());
 
     if (llvm::Value *refPtr =
             getRefPtrIfDeclareTarget(offloadPtr, moduleTranslation)) {
@@ -6688,11 +6707,11 @@ static void collectMapDataFromMapOperands(
     // field, the pointer address for the base address field, and the pointer
     // not the data (base addresses) size. So we end up with a mix of base
     // types and sizes we wish to insert here.
-    mlir::Type sizeType = (isRefPtrOrPteeMapWithAttach || !mapOp.getVarPtrPtr())
+    mlir::Type sizeType = (isAttachStyleMap || !mapOp.getVarPtrPtr())
                               ? mapOp.getVarPtrType()
                               : mapOp.getVarPtrPtrType().value();
     mapData.Sizes.push_back(getSizeInBytes(
-        dl, sizeType, isRefPtrOrPteeMapWithAttach ? nullptr : mapOp,
+        dl, sizeType, isAttachStyleMap ? nullptr : mapOp,
         mapData.Pointers.back(), moduleTranslation.convertType(sizeType),
         builder, moduleTranslation));
     mapData.MapClause.push_back(mapOp.getOperation());
@@ -7424,16 +7443,26 @@ static void processMapWithMembersOf(LLVM::ModuleTranslation &moduleTranslation,
 
   llvm::omp::OpenMPOffloadMappingFlags memberOfFlag =
       ompBuilder.getMemberOfFlag(combinedInfo.Types.size());
-  for (size_t i = 0; i < mapInfoIdx.size(); i++) {
-    // Index == 0 is the parent map and if it gets here it's an unattachable
-    // type and should have OMP_MAP_TARGET_PARAM applied and no MEMBER_OF flag.
-    if (i == 0) {
+
+  // The first index is the parent map, the rest are its members. The parent
+  // normally undergoes the standard parent-with-members mapping, contributing
+  // the MEMBER_OF flag that binds each member to it. The one exception is a
+  // privatizeable attach map (a privatized pointer/descriptor passed directly
+  // as a kernel argument): here the parent is emitted as an individual map
+  // instead, for the time being, as it's used only in pointer/allocatable to
+  // array cases for the moment. This only ever applies to the parent, so it is
+  // checked once here rather than inside the loop below.
+  bool parentIsPrivatizeableAttach =
+      isPrivatizeableAttachMap(parentClause.getMapType());
+  for (auto [i, idx] : llvm::enumerate(mapInfoIdx)) {
+    bool emitParentMap = i == 0 && !parentIsPrivatizeableAttach;
+    if (emitParentMap) {
       mapParentWithMembers(moduleTranslation, builder, ompBuilder, dl,
-                           combinedInfo, mapData, mapInfoIdx[i], memberOfFlag,
+                           combinedInfo, mapData, idx, memberOfFlag,
                            targetDirective);
     } else {
-      processIndividualMap(builder, ompBuilder, mapData, mapInfoIdx[i],
-                           combinedInfo, targetDirective, memberOfFlag,
+      processIndividualMap(builder, ompBuilder, mapData, idx, combinedInfo,
+                           targetDirective, memberOfFlag,
                            /*isTargetParam=*/false, mapDataIndex);
     }
   }
@@ -9056,17 +9085,24 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
 
   for (size_t i = 0, e = mapData.OriginalValue.size(); i != e; ++i) {
     // 1) Declare target arguments are not passed to kernels as arguments.
-    // 2) Attach maps are not passed in as arguments to kernels.
+    // 2) Attach maps are not passed in as arguments to kernels, except for
+    //    private attach maps used for corresponding-pointer initialization.
     // 3) Children of record objects are not passed in as arguments.
     // TODO: We currently do not handle cases where a member is explicitly
     // passed in as an argument, this will likley need to be handled in
     // the near future, rather than using IsAMember, it may be better to
     // test if the relevant BlockArg is used within the target region and
     // then use that as a basis for exclusion in the kernel inputs.
-    bool isAttachMap = (mapData.Types[i] &
-                        llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ATTACH) ==
-                       llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ATTACH;
-    if (!mapData.IsDeclareTarget[i] && !mapData.IsAMember[i] && !isAttachMap)
+    using MapFlags = llvm::omp::OpenMPOffloadMappingFlags;
+    bool isAttachMap = (mapData.Types[i] & MapFlags::OMP_MAP_ATTACH) ==
+                       MapFlags::OMP_MAP_ATTACH;
+    bool isPrivateTargetParam =
+        (mapData.Types[i] &
+         (MapFlags::OMP_MAP_PRIVATE | MapFlags::OMP_MAP_TARGET_PARAM)) ==
+        (MapFlags::OMP_MAP_PRIVATE | MapFlags::OMP_MAP_TARGET_PARAM);
+
+    if (!mapData.IsDeclareTarget[i] && !mapData.IsAMember[i] &&
+        (!isAttachMap || (isAttachMap && isPrivateTargetParam)))
       kernelInput.push_back(mapData.OriginalValue[i]);
   }
 
diff --git a/mlir/test/Target/LLVMIR/omptarget-map-pointer-privatization.mlir b/mlir/test/Target/LLVMIR/omptarget-map-pointer-privatization.mlir
new file mode 100644
index 0000000000000..f05355197ea88
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-map-pointer-privatization.mlir
@@ -0,0 +1,39 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// This test verifies that a privatized pointer map (a parent map carrying the
+// target_param | private | attach map type combination) is lowered such that:
+//   * the privatized parent is emitted as an individual map entry rather than
+//     undergoing the standard parent-with-members mapping, and
+//   * the parent which now has the attach map type is still passed as a kernel
+//     argument (OMP_MAP_TARGET_PARAM), unlike normal attach maps.
+
+module attributes {omp.is_gpu = false, omp.is_target_device = false, omp.requires = #omp<clause_requires none>, omp.target_triples = ["amdgcn-amd-amdhsa"], omp.version = #omp.version<version = 52>} {
+  llvm.func @assumed_shape_array_priv_(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+    %member = omp.map.info var_ptr(%arg0 : !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)>) map_clauses(tofrom) capture(ByRef) var_ptr_ptr(%arg1 : !llvm.ptr, i32) -> !llvm.ptr {name = ""}
+    %parent = omp.map.info var_ptr(%arg0 : !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)>) map_clauses(target_param, private, attach) capture(ByRef) var_ptr_ptr(%arg1 : !llvm.ptr, i32) members(%member : [0] : !llvm.ptr) -> !llvm.ptr {name = "arr_read_write"}
+    %attach = omp.map.info var_ptr(%arg0 : !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)>) map_clauses(attach, ref_ptr, ref_ptee) capture(ByRef) var_ptr_ptr(%arg1 : !llvm.ptr, i32) -> !llvm.ptr {name = "arr_read_write"}
+    omp.target kernel_type(generic) map_entries(%parent -> %arg2, %attach -> %arg3, %member -> %arg4 : !llvm.ptr, !llvm.ptr, !llvm.ptr) {
+      omp.terminator
+    }
+    llvm.return
+  }
+}
+
+// CHECK: @.offload_maptypes = private unnamed_addr constant [4 x i64] [i64 16544, i64 3, i64 16384, i64 288]
+
+// CHECK: define void @assumed_shape_array_priv_(ptr %[[ARG0:.*]], ptr %[[ARG1:.*]])
+// CHECK:  %[[MEMBER_PTR:.*]] = load ptr, ptr %[[ARG1]], align 8
+// CHECK:  %[[ATTACH_PTR:.*]] = load ptr, ptr %[[ARG1]], align 8
+// CHECK:  %[[PARENT_PTR:.*]] = load ptr, ptr %[[ARG1]], align 8
+// CHECK:  %[[BASEPTRS0:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
+// CHECK:  store ptr %[[ARG0]], ptr %[[BASEPTRS0]], align 8
+// CHECK:  %[[OFFPTRS0:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_ptrs, i32 0, i32 0
+// CHECK:  store ptr %[[MEMBER_PTR]], ptr %[[OFFPTRS0]], align 8
+// CHECK:  %[[BASEPTRS1:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_baseptrs, i32 0, i32 1
+// CHECK:  store ptr %[[ARG0]], ptr %[[BASEPTRS1]], align 8
+// CHECK:  %[[OFFPTRS1:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_ptrs, i32 0, i32 1
+// CHECK:  store ptr %[[PARENT_PTR]], ptr %[[OFFPTRS1]], align 8
+// CHECK:  %[[BASEPTRS2:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_baseptrs, i32 0, i32 2
+// CHECK:  store ptr %[[ARG0]], ptr %[[BASEPTRS2]], align 8
+// CHECK:  %[[OFFPTRS2:.*]] = getelementptr inbounds [4 x ptr], ptr %.offload_ptrs, i32 0, i32 2
+// CHECK:  store ptr %[[ATTACH_PTR]], ptr %[[OFFPTRS2]], align 8



More information about the flang-commits mailing list