shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
llvm-project/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
agozillon b291cfcad4
[Flang][OpenMP] Generate correct present checks for implicit maps of optional allocatables (#138210) 
Currently, we do not generate the appropriate checks to check if an
optional
allocatable argument is present before accessing relevant components of
it,
in particular when creating bounds, we must generate a presence check
and we
must make sure we do not generate/keep an load external to the presence
check
by utilising the raw address rather than the regular address of the info
data structure.

Similarly in cases for optional allocatables we must treat them like
non-allocatable
arguments and generate an intermediate allocation that we can have as a
location
in memory that we can access later in the lowering without causing
segfaults when
we perform "mapping" on it, even if the end result is an empty
allocatable
(basically, we shouldn't explode if someone tries to map a non-present
optional,
similar to C++ when mapping null data).
2025-05-09 13:57:45 +02:00

752 lines
34 KiB
C++
Raw Blame History

//===- MapInfoFinalization.cpp -----------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
/// \file
/// An OpenMP dialect related pass for FIR/HLFIR which performs some
/// pre-processing of MapInfoOp's after the module has been lowered to
/// finalize them.
///
/// For example, it expands MapInfoOp's containing descriptor related
/// types (fir::BoxType's) into multiple MapInfoOp's containing the parent
/// descriptor and pointer member components for individual mapping,
/// treating the descriptor type as a record type for later lowering in the
/// OpenMP dialect.
///
/// The pass also adds MapInfoOp's that are members of a parent object but are
/// not directly used in the body of a target region to its BlockArgument list
/// to maintain consistency across all MapInfoOp's tied to a region directly or
/// indirectly via a parent object.
//===----------------------------------------------------------------------===//
#include "flang/Optimizer/Builder/DirectivesCommon.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/HLFIRTools.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/KindMapping.h"
#include "flang/Optimizer/HLFIR/HLFIROps.h"
#include "flang/Optimizer/OpenMP/Passes.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Frontend/OpenMP/OMPConstants.h"
#include <algorithm>
#include <cstddef>
#include <iterator>
#include <numeric>
namespace flangomp {
#define GEN_PASS_DEF_MAPINFOFINALIZATIONPASS
#include "flang/Optimizer/OpenMP/Passes.h.inc"
} // namespace flangomp
namespace {
class MapInfoFinalizationPass
: public flangomp::impl::MapInfoFinalizationPassBase<
MapInfoFinalizationPass> {
/// Helper class tracking a members parent and its
/// placement in the parents member list
struct ParentAndPlacement {
mlir::omp::MapInfoOp parent;
size_t index;
};
/// Tracks any intermediate function/subroutine local allocations we
/// generate for the descriptors of box type dummy arguments, so that
/// we can retrieve it for subsequent reuses within the functions
/// scope.
///
/// descriptor defining op
/// | corresponding local alloca
/// | |
std::map<mlir::Operation *, mlir::Value> localBoxAllocas;
/// getMemberUserList gathers all users of a particular MapInfoOp that are
/// other MapInfoOp's and places them into the mapMemberUsers list, which
/// records the map that the current argument MapInfoOp "op" is part of
/// alongside the placement of "op" in the recorded users members list. The
/// intent of the generated list is to find all MapInfoOp's that may be
/// considered parents of the passed in "op" and in which it shows up in the
/// member list, alongside collecting the placement information of "op" in its
/// parents member list.
void
getMemberUserList(mlir::omp::MapInfoOp op,
llvm::SmallVectorImpl<ParentAndPlacement> &mapMemberUsers) {
for (auto *user : op->getUsers())
if (auto map = mlir::dyn_cast_if_present<mlir::omp::MapInfoOp>(user))
for (auto [i, mapMember] : llvm::enumerate(map.getMembers()))
if (mapMember.getDefiningOp() == op)
mapMemberUsers.push_back({map, i});
}
void getAsIntegers(llvm::ArrayRef<mlir::Attribute> values,
llvm::SmallVectorImpl<int64_t> &ints) {
ints.reserve(values.size());
llvm::transform(values, std::back_inserter(ints),
[](mlir::Attribute value) {
return mlir::cast<mlir::IntegerAttr>(value).getInt();
});
}
/// This function will expand a MapInfoOp's member indices back into a vector
/// so that they can be trivially modified as unfortunately the attribute type
/// that's used does not have modifiable fields at the moment (generally
/// awkward to work with)
void getMemberIndicesAsVectors(
mlir::omp::MapInfoOp mapInfo,
llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &indices) {
indices.reserve(mapInfo.getMembersIndexAttr().getValue().size());
llvm::transform(mapInfo.getMembersIndexAttr().getValue(),
std::back_inserter(indices), [this](mlir::Attribute value) {
auto memberIndex = mlir::cast<mlir::ArrayAttr>(value);
llvm::SmallVector<int64_t> indexes;
getAsIntegers(memberIndex.getValue(), indexes);
return indexes;
});
}
/// When provided a MapInfoOp containing a descriptor type that
/// we must expand into multiple maps this function will extract
/// the value from it and return it, in certain cases we must
/// generate a new allocation to store into so that the
/// fir::BoxOffsetOp we utilise to access the descriptor datas
/// base address can be utilised.
mlir::Value getDescriptorFromBoxMap(mlir::omp::MapInfoOp boxMap,
fir::FirOpBuilder &builder) {
mlir::Value descriptor = boxMap.getVarPtr();
if (!fir::isTypeWithDescriptor(boxMap.getVarType()))
if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
boxMap.getVarPtr().getDefiningOp()))
descriptor = addrOp.getVal();
if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
!fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
return descriptor;
mlir::Value &slot = localBoxAllocas[descriptor.getDefiningOp()];
if (slot) {
return slot;
}
// The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as
// allowing it to access non-reference box operations can cause some
// problematic SSA IR. However, in the case of assumed shape's the type
// is not a !fir.ref, in these cases to retrieve the appropriate
// !fir.ref<!fir.box<...>> to access the data we need to map we must
// perform an alloca and then store to it and retrieve the data from the new
// alloca.
mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint();
mlir::Block *allocaBlock = builder.getAllocaBlock();
mlir::Location loc = boxMap->getLoc();
assert(allocaBlock && "No alloca block found for this top level op");
builder.setInsertionPointToStart(allocaBlock);
mlir::Type allocaType = descriptor.getType();
if (fir::isBoxAddress(allocaType))
allocaType = fir::unwrapRefType(allocaType);
auto alloca = builder.create<fir::AllocaOp>(loc, allocaType);
builder.restoreInsertionPoint(insPt);
// We should only emit a store if the passed in data is present, it is
// possible a user passes in no argument to an optional parameter, in which
// case we cannot store or we'll segfault on the emitted memcpy.
auto isPresent =
builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), descriptor);
builder.genIfOp(loc, {}, isPresent, false)
.genThen([&]() {
descriptor = builder.loadIfRef(loc, descriptor);
builder.create<fir::StoreOp>(loc, descriptor, alloca);
})
.end();
return slot = alloca;
}
/// Function that generates a FIR operation accessing the descriptor's
/// base address (BoxOffsetOp) and a MapInfoOp for it. The most
/// important thing to note is that we normally move the bounds from
/// the descriptor map onto the base address map.
mlir::omp::MapInfoOp genBaseAddrMap(mlir::Value descriptor,
mlir::OperandRange bounds,
int64_t mapType,
fir::FirOpBuilder &builder) {
mlir::Location loc = descriptor.getLoc();
mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
loc, descriptor, fir::BoxFieldAttr::base_addr);
mlir::Type underlyingVarType =
llvm::cast<mlir::omp::PointerLikeType>(
fir::unwrapRefType(baseAddrAddr.getType()))
.getElementType();
if (auto seqType = llvm::dyn_cast<fir::SequenceType>(underlyingVarType))
if (seqType.hasDynamicExtents())
underlyingVarType = seqType.getEleTy();
// Member of the descriptor pointing at the allocated data
return builder.create<mlir::omp::MapInfoOp>(
loc, baseAddrAddr.getType(), descriptor,
mlir::TypeAttr::get(underlyingVarType),
builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
mlir::omp::VariableCaptureKind::ByRef),
baseAddrAddr, /*members=*/mlir::SmallVector<mlir::Value>{},
/*membersIndex=*/mlir::ArrayAttr{}, bounds,
/*mapperId*/ mlir::FlatSymbolRefAttr(),
/*name=*/builder.getStringAttr(""),
/*partial_map=*/builder.getBoolAttr(false));
}
/// This function adjusts the member indices vector to include a new
/// base address member. We take the position of the descriptor in
/// the member indices list, which is the index data that the base
/// addresses index will be based off of, as the base address is
/// a member of the descriptor. We must also alter other members
/// that are members of this descriptor to account for the addition
/// of the base address index.
void adjustMemberIndices(
llvm::SmallVectorImpl<llvm::SmallVector<int64_t>> &memberIndices,
size_t memberIndex) {
llvm::SmallVector<int64_t> baseAddrIndex = memberIndices[memberIndex];
// If we find another member that is "derived/a member of" the descriptor
// that is not the descriptor itself, we must insert a 0 for the new base
// address we have just added for the descriptor into the list at the
// appropriate position to maintain correctness of the positional/index data
// for that member.
for (llvm::SmallVector<int64_t> &member : memberIndices)
if (member.size() > baseAddrIndex.size() &&
std::equal(baseAddrIndex.begin(), baseAddrIndex.end(),
member.begin()))
member.insert(std::next(member.begin(), baseAddrIndex.size()), 0);
// Add the base address index to the main base address member data
baseAddrIndex.push_back(0);
// Insert our newly created baseAddrIndex into the larger list of indices at
// the correct location.
memberIndices.insert(std::next(memberIndices.begin(), memberIndex + 1),
baseAddrIndex);
}
/// Adjusts the descriptor's map type. The main alteration that is done
/// currently is transforming the map type to `OMP_MAP_TO` where possible.
/// This is because we will always need to map the descriptor to device
/// (or at the very least it seems to be the case currently with the
/// current lowered kernel IR), as without the appropriate descriptor
/// information on the device there is a risk of the kernel IR
/// requesting for various data that will not have been copied to
/// perform things like indexing. This can cause segfaults and
/// memory access errors. However, we do not need this data mapped
/// back to the host from the device, as per the OpenMP spec we cannot alter
/// the data via resizing or deletion on the device. Discarding any
/// descriptor alterations via no map back is reasonable (and required
/// for certain segments of descriptor data like the type descriptor that are
/// global constants). This alteration is only inapplicable to `target exit`
/// and `target update` currently, and that's due to `target exit` not
/// allowing `to` mappings, and `target update` not allowing both `to` and
/// `from` simultaneously. We currently try to maintain the `implicit` flag
/// where necessary, although it does not seem strictly required.
unsigned long getDescriptorMapType(unsigned long mapTypeFlag,
mlir::Operation *target) {
using mapFlags = llvm::omp::OpenMPOffloadMappingFlags;
if (llvm::isa_and_nonnull<mlir::omp::TargetExitDataOp,
mlir::omp::TargetUpdateOp>(target))
return mapTypeFlag;
mapFlags flags = mapFlags::OMP_MAP_TO |
(mapFlags(mapTypeFlag) &
(mapFlags::OMP_MAP_IMPLICIT | mapFlags::OMP_MAP_CLOSE |
mapFlags::OMP_MAP_ALWAYS));
return llvm::to_underlying(flags);
}
/// Check if the mapOp is present in the HasDeviceAddr clause on
/// the userOp. Only applies to TargetOp.
bool isHasDeviceAddr(mlir::omp::MapInfoOp mapOp, mlir::Operation *userOp) {
assert(userOp && "Expecting non-null argument");
if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(userOp)) {
for (mlir::Value hda : targetOp.getHasDeviceAddrVars()) {
if (hda.getDefiningOp() == mapOp)
return true;
}
}
return false;
}
mlir::omp::MapInfoOp genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
fir::FirOpBuilder &builder,
mlir::Operation *target) {
llvm::SmallVector<ParentAndPlacement> mapMemberUsers;
getMemberUserList(op, mapMemberUsers);
// TODO: map the addendum segment of the descriptor, similarly to the
// base address/data pointer member.
mlir::Value descriptor = getDescriptorFromBoxMap(op, builder);
mlir::ArrayAttr newMembersAttr;
mlir::SmallVector<mlir::Value> newMembers;
llvm::SmallVector<llvm::SmallVector<int64_t>> memberIndices;
bool IsHasDeviceAddr = isHasDeviceAddr(op, target);
if (!mapMemberUsers.empty() || !op.getMembers().empty())
getMemberIndicesAsVectors(
!mapMemberUsers.empty() ? mapMemberUsers[0].parent : op,
memberIndices);
// If the operation that we are expanding with a descriptor has a user
// (parent), then we have to expand the parent's member indices to reflect
// the adjusted member indices for the base address insertion. However, if
// it does not then we are expanding a MapInfoOp without any pre-existing
// member information to now have one new member for the base address, or
// we are expanding a parent that is a descriptor and we have to adjust
// all of its members to reflect the insertion of the base address.
//
// If we're expanding a top-level descriptor for a map operation that
// resulted from "has_device_addr" clause, then we want the base pointer
// from the descriptor to be used verbatim, i.e. without additional
// remapping. To avoid this remapping, simply don't generate any map
// information for the descriptor members.
if (!mapMemberUsers.empty()) {
// Currently, there should only be one user per map when this pass
// is executed. Either a parent map, holding the current map in its
// member list, or a target operation that holds a map clause. This
// may change in the future if we aim to refactor the MLIR for map
// clauses to allow sharing of duplicate maps across target
// operations.
assert(mapMemberUsers.size() == 1 &&
"OMPMapInfoFinalization currently only supports single users of a "
"MapInfoOp");
auto baseAddr =
genBaseAddrMap(descriptor, op.getBounds(), op.getMapType(), builder);
ParentAndPlacement mapUser = mapMemberUsers[0];
adjustMemberIndices(memberIndices, mapUser.index);
llvm::SmallVector<mlir::Value> newMemberOps;
for (auto v : mapUser.parent.getMembers()) {
newMemberOps.push_back(v);
if (v == op)
newMemberOps.push_back(baseAddr);
}
mapUser.parent.getMembersMutable().assign(newMemberOps);
mapUser.parent.setMembersIndexAttr(
builder.create2DI64ArrayAttr(memberIndices));
} else if (!IsHasDeviceAddr) {
auto baseAddr =
genBaseAddrMap(descriptor, op.getBounds(), op.getMapType(), builder);
newMembers.push_back(baseAddr);
if (!op.getMembers().empty()) {
for (auto &indices : memberIndices)
indices.insert(indices.begin(), 0);
memberIndices.insert(memberIndices.begin(), {0});
newMembersAttr = builder.create2DI64ArrayAttr(memberIndices);
newMembers.append(op.getMembers().begin(), op.getMembers().end());
} else {
llvm::SmallVector<llvm::SmallVector<int64_t>> memberIdx = {{0}};
newMembersAttr = builder.create2DI64ArrayAttr(memberIdx);
}
}
// Descriptors for objects listed on the `has_device_addr` will always
// be copied. This is because the descriptor can be rematerialized by the
// compiler, and so the address of the descriptor for a given object at
// one place in the code may differ from that address in another place.
// The contents of the descriptor (the base address in particular) will
// remain unchanged though.
uint64_t mapType = op.getMapType();
if (IsHasDeviceAddr) {
mapType |= llvm::to_underlying(
llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_ALWAYS);
}
mlir::omp::MapInfoOp newDescParentMapOp =
builder.create<mlir::omp::MapInfoOp>(
op->getLoc(), op.getResult().getType(), descriptor,
mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
builder.getIntegerAttr(builder.getIntegerType(64, false),
getDescriptorMapType(mapType, target)),
op.getMapCaptureTypeAttr(), /*varPtrPtr=*/mlir::Value{}, newMembers,
newMembersAttr, /*bounds=*/mlir::SmallVector<mlir::Value>{},
/*mapperId*/ mlir::FlatSymbolRefAttr(), op.getNameAttr(),
/*partial_map=*/builder.getBoolAttr(false));
op.replaceAllUsesWith(newDescParentMapOp.getResult());
op->erase();
return newDescParentMapOp;
}
// We add all mapped record members not directly used in the target region
// to the block arguments in front of their parent and we place them into
// the map operands list for consistency.
//
// These indirect uses (via accesses to their parent) will still be
// mapped individually in most cases, and a parent mapping doesn't
// guarantee the parent will be mapped in its totality, partial
// mapping is common.
//
// For example:
// map(tofrom: x%y)
//
// Will generate a mapping for "x" (the parent) and "y" (the member).
// The parent "x" will not be mapped, but the member "y" will.
// However, we must have the parent as a BlockArg and MapOperand
// in these cases, to maintain the correct uses within the region and
// to help tracking that the member is part of a larger object.
//
// In the case of:
// map(tofrom: x%y, x%z)
//
// The parent member becomes more critical, as we perform a partial
// structure mapping where we link the mapping of the members y
// and z together via the parent x. We do this at a kernel argument
// level in LLVM IR and not just MLIR, which is important to maintain
// similarity to Clang and for the runtime to do the correct thing.
// However, we still do not map the structure in its totality but
// rather we generate an un-sized "binding" map entry for it.
//
// In the case of:
// map(tofrom: x, x%y, x%z)
//
// We do actually map the entirety of "x", so the explicit mapping of
// x%y, x%z becomes unnecessary. It is redundant to write this from a
// Fortran OpenMP perspective (although it is legal), as even if the
// members were allocatables or pointers, we are mandated by the
// specification to map these (and any recursive components) in their
// entirety, which is different to the C++ equivalent, which requires
// explicit mapping of these segments.
void addImplicitMembersToTarget(mlir::omp::MapInfoOp op,
fir::FirOpBuilder &builder,
mlir::Operation *target) {
auto mapClauseOwner =
llvm::dyn_cast_if_present<mlir::omp::MapClauseOwningOpInterface>(
target);
// TargetDataOp is technically a MapClauseOwningOpInterface, so we
// do not need to explicitly check for the extra cases here for use_device
// addr/ptr
if (!mapClauseOwner)
return;
auto addOperands = [&](mlir::MutableOperandRange &mutableOpRange,
mlir::Operation *directiveOp,
unsigned blockArgInsertIndex = 0) {
if (!llvm::is_contained(mutableOpRange.getAsOperandRange(),
op.getResult()))
return;
// There doesn't appear to be a simple way to convert MutableOperandRange
// to a vector currently, so we instead use a for_each to populate our
// vector.
llvm::SmallVector<mlir::Value> newMapOps;
newMapOps.reserve(mutableOpRange.size());
llvm::for_each(
mutableOpRange.getAsOperandRange(),
[&newMapOps](mlir::Value oper) { newMapOps.push_back(oper); });
for (auto mapMember : op.getMembers()) {
if (llvm::is_contained(mutableOpRange.getAsOperandRange(), mapMember))
continue;
newMapOps.push_back(mapMember);
if (directiveOp) {
directiveOp->getRegion(0).insertArgument(
blockArgInsertIndex, mapMember.getType(), mapMember.getLoc());
blockArgInsertIndex++;
}
}
mutableOpRange.assign(newMapOps);
};
auto argIface =
llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(target);
if (auto mapClauseOwner =
llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
mlir::MutableOperandRange mapMutableOpRange =
mapClauseOwner.getMapVarsMutable();
unsigned blockArgInsertIndex =
argIface
? argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs()
: 0;
addOperands(mapMutableOpRange,
llvm::dyn_cast_if_present<mlir::omp::TargetOp>(
argIface.getOperation()),
blockArgInsertIndex);
}
if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(target)) {
mlir::MutableOperandRange useDevAddrMutableOpRange =
targetDataOp.getUseDeviceAddrVarsMutable();
addOperands(useDevAddrMutableOpRange, target,
argIface.getUseDeviceAddrBlockArgsStart() +
argIface.numUseDeviceAddrBlockArgs());
mlir::MutableOperandRange useDevPtrMutableOpRange =
targetDataOp.getUseDevicePtrVarsMutable();
addOperands(useDevPtrMutableOpRange, target,
argIface.getUseDevicePtrBlockArgsStart() +
argIface.numUseDevicePtrBlockArgs());
} else if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target)) {
mlir::MutableOperandRange hasDevAddrMutableOpRange =
targetOp.getHasDeviceAddrVarsMutable();
addOperands(hasDevAddrMutableOpRange, target,
argIface.getHasDeviceAddrBlockArgsStart() +
argIface.numHasDeviceAddrBlockArgs());
}
}
// We retrieve the first user that is a Target operation, of which
// there should only be one currently. Every MapInfoOp can be tied to
// at most one Target operation and at the minimum no operations.
// This may change in the future with IR cleanups/modifications,
// in which case this pass will need updating to support cases
// where a map can have more than one user and more than one of
// those users can be a Target operation. For now, we simply
// return the first target operation encountered, which may
// be on the parent MapInfoOp in the case of a member mapping.
// In that case, we traverse the MapInfoOp chain until we
// find the first TargetOp user.
mlir::Operation *getFirstTargetUser(mlir::omp::MapInfoOp mapOp) {
for (auto *user : mapOp->getUsers()) {
if (llvm::isa<mlir::omp::TargetOp, mlir::omp::TargetDataOp,
mlir::omp::TargetUpdateOp, mlir::omp::TargetExitDataOp,
mlir::omp::TargetEnterDataOp,
mlir::omp::DeclareMapperInfoOp>(user))
return user;
if (auto mapUser = llvm::dyn_cast<mlir::omp::MapInfoOp>(user))
return getFirstTargetUser(mapUser);
}
return nullptr;
}
// This pass executes on omp::MapInfoOp's containing descriptor based types
// (allocatables, pointers, assumed shape etc.) and expanding them into
// multiple omp::MapInfoOp's for each pointer member contained within the
// descriptor.
//
// From the perspective of the MLIR pass manager this runs on the top level
// operation (usually function) containing the MapInfoOp because this pass
// will mutate siblings of MapInfoOp.
void runOnOperation() override {
mlir::ModuleOp module = getOperation();
if (!module)
module = getOperation()->getParentOfType<mlir::ModuleOp>();
fir::KindMapping kindMap = fir::getKindMapping(module);
fir::FirOpBuilder builder{module, std::move(kindMap)};
// We wish to maintain some function level scope (currently
// just local function scope variables used to load and store box
// variables into so we can access their base address, an
// quirk of box_offset requires us to have an in memory box, but Fortran
// in certain cases does not provide this) whilst not subjecting
// ourselves to the possibility of race conditions while this pass
// undergoes frequent re-iteration for the near future. So we loop
// over function in the module and then map.info inside of those.
getOperation()->walk([&](mlir::Operation *func) {
if (!mlir::isa<mlir::func::FuncOp, mlir::omp::DeclareMapperOp>(func))
return;
// clear all local allocations we made for any boxes in any prior
// iterations from previous function scopes.
localBoxAllocas.clear();
// First, walk `omp.map.info` ops to see if any record members should be
// implicitly mapped.
func->walk([&](mlir::omp::MapInfoOp op) {
mlir::Type underlyingType =
fir::unwrapRefType(op.getVarPtr().getType());
// TODO Test with and support more complicated cases; like arrays for
// records, for example.
if (!fir::isRecordWithAllocatableMember(underlyingType))
return mlir::WalkResult::advance();
// TODO For now, only consider `omp.target` ops. Other ops that support
// `map` clauses will follow later.
mlir::omp::TargetOp target =
mlir::dyn_cast_if_present<mlir::omp::TargetOp>(
getFirstTargetUser(op));
if (!target)
return mlir::WalkResult::advance();
auto mapClauseOwner =
llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(*target);
int64_t mapVarIdx = mapClauseOwner.getOperandIndexForMap(op);
assert(mapVarIdx >= 0 &&
mapVarIdx <
static_cast<int64_t>(mapClauseOwner.getMapVars().size()));
auto argIface =
llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(*target);
// TODO How should `map` block argument that correspond to: `private`,
// `use_device_addr`, `use_device_ptr`, be handled?
mlir::BlockArgument opBlockArg = argIface.getMapBlockArgs()[mapVarIdx];
llvm::SetVector<mlir::Operation *> mapVarForwardSlice;
mlir::getForwardSlice(opBlockArg, &mapVarForwardSlice);
mapVarForwardSlice.remove_if([&](mlir::Operation *sliceOp) {
// TODO Support coordinate_of ops.
//
// TODO Support call ops by recursively examining the forward slice of
// the corresponding parameter to the field in the called function.
return !mlir::isa<hlfir::DesignateOp>(sliceOp);
});
auto recordType = mlir::cast<fir::RecordType>(underlyingType);
llvm::SmallVector<mlir::Value> newMapOpsForFields;
llvm::SmallVector<int64_t> fieldIndicies;
for (auto fieldMemTyPair : recordType.getTypeList()) {
auto &field = fieldMemTyPair.first;
auto memTy = fieldMemTyPair.second;
bool shouldMapField =
llvm::find_if(mapVarForwardSlice, [&](mlir::Operation *sliceOp) {
if (!fir::isAllocatableType(memTy))
return false;
auto designateOp = mlir::dyn_cast<hlfir::DesignateOp>(sliceOp);
if (!designateOp)
return false;
return designateOp.getComponent() &&
designateOp.getComponent()->strref() == field;
}) != mapVarForwardSlice.end();
// TODO Handle recursive record types. Adapting
// `createParentSymAndGenIntermediateMaps` to work direclty on MLIR
// entities might be helpful here.
if (!shouldMapField)
continue;
int32_t fieldIdx = recordType.getFieldIndex(field);
bool alreadyMapped = [&]() {
if (op.getMembersIndexAttr())
for (auto indexList : op.getMembersIndexAttr()) {
auto indexListAttr = mlir::cast<mlir::ArrayAttr>(indexList);
if (indexListAttr.size() == 1 &&
mlir::cast<mlir::IntegerAttr>(indexListAttr[0]).getInt() ==
fieldIdx)
return true;
}
return false;
}();
if (alreadyMapped)
continue;
builder.setInsertionPoint(op);
fir::IntOrValue idxConst =
mlir::IntegerAttr::get(builder.getI32Type(), fieldIdx);
auto fieldCoord = builder.create<fir::CoordinateOp>(
op.getLoc(), builder.getRefType(memTy), op.getVarPtr(),
llvm::SmallVector<fir::IntOrValue, 1>{idxConst});
fir::factory::AddrAndBoundsInfo info =
fir::factory::getDataOperandBaseAddr(
builder, fieldCoord, /*isOptional=*/false, op.getLoc());
llvm::SmallVector<mlir::Value> bounds =
fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
mlir::omp::MapBoundsType>(
builder, info,
hlfir::translateToExtendedValue(op.getLoc(), builder,
hlfir::Entity{fieldCoord})
.first,
/*dataExvIsAssumedSize=*/false, op.getLoc());
mlir::omp::MapInfoOp fieldMapOp =
builder.create<mlir::omp::MapInfoOp>(
op.getLoc(), fieldCoord.getResult().getType(),
fieldCoord.getResult(),
mlir::TypeAttr::get(
fir::unwrapRefType(fieldCoord.getResult().getType())),
op.getMapTypeAttr(),
builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
mlir::omp::VariableCaptureKind::ByRef),
/*varPtrPtr=*/mlir::Value{}, /*members=*/mlir::ValueRange{},
/*members_index=*/mlir::ArrayAttr{}, bounds,
/*mapperId=*/mlir::FlatSymbolRefAttr(),
builder.getStringAttr(op.getNameAttr().strref() + "." +
field + ".implicit_map"),
/*partial_map=*/builder.getBoolAttr(false));
newMapOpsForFields.emplace_back(fieldMapOp);
fieldIndicies.emplace_back(fieldIdx);
}
if (newMapOpsForFields.empty())
return mlir::WalkResult::advance();
op.getMembersMutable().append(newMapOpsForFields);
llvm::SmallVector<llvm::SmallVector<int64_t>> newMemberIndices;
mlir::ArrayAttr oldMembersIdxAttr = op.getMembersIndexAttr();
if (oldMembersIdxAttr)
for (mlir::Attribute indexList : oldMembersIdxAttr) {
llvm::SmallVector<int64_t> listVec;
for (mlir::Attribute index : mlir::cast<mlir::ArrayAttr>(indexList))
listVec.push_back(mlir::cast<mlir::IntegerAttr>(index).getInt());
newMemberIndices.emplace_back(std::move(listVec));
}
for (int64_t newFieldIdx : fieldIndicies)
newMemberIndices.emplace_back(
llvm::SmallVector<int64_t>(1, newFieldIdx));
op.setMembersIndexAttr(builder.create2DI64ArrayAttr(newMemberIndices));
op.setPartialMap(true);
return mlir::WalkResult::advance();
});
func->walk([&](mlir::omp::MapInfoOp op) {
// TODO: Currently only supports a single user for the MapInfoOp. This
// is fine for the moment, as the Fortran frontend will generate a
// new MapInfoOp with at most one user currently. In the case of
// members of other objects, like derived types, the user would be the
// parent. In cases where it's a regular non-member map, the user would
// be the target operation it is being mapped by.
//
// However, when/if we optimise/cleanup the IR we will have to extend
// this pass to support multiple users, as we may wish to have a map
// be re-used by multiple users (e.g. across multiple targets that map
// the variable and have identical map properties).
assert(llvm::hasSingleElement(op->getUsers()) &&
"OMPMapInfoFinalization currently only supports single users "
"of a MapInfoOp");
if (fir::isTypeWithDescriptor(op.getVarType()) ||
mlir::isa_and_present<fir::BoxAddrOp>(
op.getVarPtr().getDefiningOp())) {
builder.setInsertionPoint(op);
mlir::Operation *targetUser = getFirstTargetUser(op);
assert(targetUser && "expected user of map operation was not found");
genDescriptorMemberMaps(op, builder, targetUser);
}
});
// Wait until after we have generated all of our maps to add them onto
// the target's block arguments, simplifying the process as there would be
// no need to avoid accidental duplicate additions.
func->walk([&](mlir::omp::MapInfoOp op) {
mlir::Operation *targetUser = getFirstTargetUser(op);
assert(targetUser && "expected user of map operation was not found");
addImplicitMembersToTarget(op, builder, targetUser);
});
});
}
};
} // namespace
Reference in New Issue View Git Blame Copy Permalink