[flang-commits] [flang] ac2390a - [Flang][OpenMP] Privatize descriptors for assumed shape array maps for performance increase (#212336)
via flang-commits
flang-commits at lists.llvm.org
Thu Aug 27 11:56:40 PDT 2026
Author: agozillon
Date: 2026-08-27T20:56:34+02:00
New Revision: ac2390a85ea0d71f3d4ed63b9cb4d68e199dc933
URL: https://github.com/llvm/llvm-project/commit/ac2390a85ea0d71f3d4ed63b9cb4d68e199dc933
DIFF: https://github.com/llvm/llvm-project/commit/ac2390a85ea0d71f3d4ed63b9cb4d68e199dc933.diff
LOG: [Flang][OpenMP] Privatize descriptors for assumed shape array maps for performance increase (#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>
Added:
flang/test/Lower/OpenMP/map-descriptor-privatization.f90
mlir/test/Target/LLVMIR/omptarget-map-pointer-privatization.mlir
Modified:
flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
flang/test/Lower/OpenMP/array-bounds.f90
flang/test/Lower/OpenMP/map-descriptor-deferral.f90
flang/test/Lower/OpenMP/optional-argument-map-3.f90
mlir/include/mlir/Dialect/OpenMP/OpenMPEnums.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
Removed:
################################################################################
diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
index 8fc3d85ce7a7e..0fdc2e1278589 100644
--- a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
+++ b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -419,14 +419,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()))
@@ -831,10 +840,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;
@@ -890,6 +919,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
@@ -911,12 +952,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(),
@@ -1141,7 +1178,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) &&
@@ -1168,19 +1206,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{};
+ bool baseAddrInsert = optDescMap && baseAddr;
+ if (baseAddrInsert)
+ 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(),
+ baseAddrInsert ? baseAddr.getVarPtrPtr() : mlir::Value{},
+ baseAddrInsert ? mlir::TypeAttr::get(underlyingVarType)
+ : mlir::TypeAttr{},
+ newMembers, newMembersAttr,
/*bounds=*/mlir::SmallVector<mlir::Value>{},
/*mapperId*/ mlir::FlatSymbolRefAttr(), op.getNameAttr(),
/*partial_map=*/builder.getBoolAttr(false));
@@ -1216,6 +1265,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);
@@ -1235,10 +1285,16 @@ 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();
+ // Exclude irregular maps from optimization via privatization; at least for
+ // the moment.
+ if (isHasDeviceAddrFlag || isUseDeviceAddr(op, *target) ||
+ isUseDevicePtr(op, *target))
+ canOptimizeDescViaPrivatization = false;
+
// If we're a derived type descriptor, that's been flagged as ref_ptr,
// but, in the same mapping, we also have members with their own
// descriptors also mapped as ref_ptr, then we have to map the parent
@@ -1262,9 +1318,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 f3a0a14dc4459..4a338fe7906a7 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]]) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
-!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>>>) name("arr_read_write(2:5)") -> !fir.ref<!fir.array<?xi32>>
+!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>>>) name("arr_read_write(2:5)") -> !fir.ref<!fir.array<?xi32>>
!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]]) name("arr_read_write(2:5)") -> !fir.ref<!fir.array<?xi32>>
!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 8e8c575ab8e85..72e084a29dbea 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(%{{.*}}) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
-!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>>>) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
+!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>>>) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
!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(%{{.*}}) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
!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(%{{.*}}) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
-!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>>>) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
+!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>>>) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
!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(%{{.*}}) name("assumed_arr") -> !fir.ref<!fir.array<?xi32>>
!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..8c5f0dd8a1890
--- /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(%{{.*}}) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xi32>>>
+!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>>>) name("arr_read_write") -> !fir.ref<!fir.array<?xi32>>
+!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(%{{.*}}) name("arr_read_write") -> !fir.ref<!fir.array<?xi32>>
+!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 1879119446f4f..e0b762f615d06 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(%{{.*}}) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>
-! 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>>>) name("dt") -> !fir.ref<!fir.array<?xf32>>
+! 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>>>) name("dt") -> !fir.ref<!fir.array<?xf32>>
! 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(%{{.*}}) name("") -> !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>
-! 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>>>) name("dt") -> !fir.ref<!fir.array<?xf32>>
+! 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>>>) name("dt") -> !fir.ref<!fir.array<?xf32>>
! 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 85ff885f4a8f3..1f210ef60fe39 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2340,6 +2340,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();
};
@@ -2371,6 +2374,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 f5e579eaa023e..e23cda41ded81 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -7087,6 +7087,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
@@ -7239,6 +7249,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)
@@ -7283,16 +7296,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)) {
@@ -7318,11 +7331,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());
@@ -8072,17 +8085,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);
}
}
}
@@ -9634,17 +9660,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..abbb0efa600aa
--- /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) name("") -> !llvm.ptr
+ %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) name("arr_read_write") -> !llvm.ptr
+ %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) name("arr_read_write") -> !llvm.ptr
+ 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