[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
Wed Aug 12 08:34:37 PDT 2026
https://github.com/agozillon updated https://github.com/llvm/llvm-project/pull/212336
>From fbff2b599ef21833cc8a1e96fb4b0eb6ef5b10a3 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 | 85 +++++++++++++++----
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 | 75 +++++++++++-----
.../omptarget-map-pointer-privatization.mlir | 39 +++++++++
9 files changed, 193 insertions(+), 44 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 949da8f20cbbe..4e041cbf5dc28 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -340,14 +340,23 @@ 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>(
boxMap.getVarPtr().getDefiningOp()))
descriptor = addrOp.getVal();
+ // We defer descriptor mapping until target or target data regions for
+ // non-allocatable, non-pointer type dummy arguments with assumed type or
+ // shape. We choose to optimize via privatization a subset of these cases
+ // for target regions, which can not be deferred. We can extend the
+ // privatization to allocatables pointers and other descriptor types as
+ // needed in the future.
canDescBeDeferred = canDeferDescriptorMapping(descriptor);
+ canOptimizeDescViaPrivatization = isDummyArgument(descriptor) &&
+ fir::isAssumedShape(descriptor.getType());
if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
!fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
@@ -752,10 +761,30 @@ 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
+ //
+ // This map type triggers the runtime to perform firstprivatization
+ // on the descriptor, treating the descriptor as a privatized entity
+ // for the duration of the device kernel, initialized with the same
+ // data as the host descriptor. The transferred data is then attached
+ // to the descriptor. The effects of this, other than the descriptor
+ // being privatized, are that the descriptors contents gets transferred
+ // across to the device with the initial kernel payload, packaged
+ // alongside the initial kernel argument list, reducing the number of
+ // host to device transfers required alongside runtime overhead as we
+ // batch as much of our required data together as we can.
+ 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;
@@ -816,6 +845,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
@@ -837,12 +878,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(),
@@ -1067,7 +1104,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) &&
@@ -1094,19 +1132,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));
@@ -1142,6 +1191,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);
@@ -1161,8 +1211,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,
@@ -1188,9 +1238,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 4a9444ca40a30..05212f5dc4de9 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2298,6 +2298,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();
};
@@ -2329,6 +2332,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 6f692e0c14db2..6a536aebbdd84 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6573,6 +6573,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
@@ -6725,6 +6735,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)
@@ -6769,16 +6782,16 @@ static void collectMapDataFromMapOperands(
// 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()) ||
+ isPrivatizeableAttachMap(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)) {
@@ -6804,11 +6817,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());
@@ -7558,17 +7571,30 @@ 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,
- /*isTargetParam=*/false, mapDataIndex);
+ processIndividualMap(
+ builder, ompBuilder, mapData, idx, combinedInfo, targetDirective,
+ parentIsPrivatizeableAttach
+ ? llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_NONE
+ : memberOfFlag,
+ /*isTargetParam=*/false, mapDataIndex);
}
}
}
@@ -9120,17 +9146,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