435e850ba9
This patch is one in a series of four patches that seeks to refactor
slightly and extend the current record type map support that was
put in place for Fortran's descriptor types to handle explicit
member mapping for record types at a single level of depth.
For example, the below case where two members of a Fortran
derived type are mapped explicitly:
''''
type :: scalar_and_array
real(4) :: real
integer(4) :: array(10)
integer(4) :: int
end type scalar_and_array
type(scalar_and_array) :: scalar_arr
!$omp target map(tofrom: scalar_arr%int, scalar_arr%real)
''''
Current cases of derived type mapping left for future work are:
> explicit member mapping of nested members (e.g. two layers of
record types where we explicitly map a member from the internal
record type)
> Fortran's automagical mapping of all elements and nested elements
of a derived type
> explicit member mapping of a derived type and then constituient members
(redundant in Fortran due to former case but still legal as far as I am aware)
> explicit member mapping of a record type (may be handled reasonably, just
not fully tested in this iteration)
> explicit member mapping for Fortran allocatable types (a variation of nested
record types)
This patch seeks to support this by extending the Flang-new OpenMP lowering to
support generation of this newly required information, creating the neccessary
parent <-to-> member map_info links, calculating the member indices and
setting if it's a partial map.
The OMPDescriptorMapInfoGen pass has also been generalized into a map
finalization phase, now named OMPMapInfoFinalization. This pass was extended
to support the insertion of member maps into the BlockArg and MapOperands of
relevant map carrying operations. Similar to the method in which descriptor types
are expanded and constituient members inserted.
Pull Request: https://github.com/llvm/llvm-project/pull/82853
262 lines
12 KiB
C++
262 lines
12 KiB
C++
//===- OMPMapInfoFinalization.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/FIRBuilder.h"
|
|
#include "flang/Optimizer/Dialect/FIRType.h"
|
|
#include "flang/Optimizer/Dialect/Support/KindMapping.h"
|
|
#include "flang/Optimizer/Transforms/Passes.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 <iterator>
|
|
|
|
namespace fir {
|
|
#define GEN_PASS_DEF_OMPMAPINFOFINALIZATIONPASS
|
|
#include "flang/Optimizer/Transforms/Passes.h.inc"
|
|
} // namespace fir
|
|
|
|
namespace {
|
|
class OMPMapInfoFinalizationPass
|
|
: public fir::impl::OMPMapInfoFinalizationPassBase<
|
|
OMPMapInfoFinalizationPass> {
|
|
|
|
void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
|
|
fir::FirOpBuilder &builder,
|
|
mlir::Operation *target) {
|
|
mlir::Location loc = op.getLoc();
|
|
mlir::Value descriptor = op.getVarPtr();
|
|
|
|
// If we enter this function, but the mapped type itself is not the
|
|
// descriptor, then it's likely the address of the descriptor so we
|
|
// must retrieve the descriptor SSA.
|
|
if (!fir::isTypeWithDescriptor(op.getVarType())) {
|
|
if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
|
|
op.getVarPtr().getDefiningOp())) {
|
|
descriptor = addrOp.getVal();
|
|
}
|
|
}
|
|
|
|
// 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.
|
|
if (mlir::isa<fir::BaseBoxType>(descriptor.getType())) {
|
|
mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint();
|
|
builder.setInsertionPointToStart(builder.getAllocaBlock());
|
|
auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
|
|
builder.restoreInsertionPoint(insPt);
|
|
builder.create<fir::StoreOp>(loc, descriptor, alloca);
|
|
descriptor = alloca;
|
|
}
|
|
|
|
mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
|
|
loc, descriptor, fir::BoxFieldAttr::base_addr);
|
|
|
|
// Member of the descriptor pointing at the allocated data
|
|
mlir::Value baseAddr = builder.create<mlir::omp::MapInfoOp>(
|
|
loc, baseAddrAddr.getType(), descriptor,
|
|
mlir::TypeAttr::get(llvm::cast<mlir::omp::PointerLikeType>(
|
|
fir::unwrapRefType(baseAddrAddr.getType()))
|
|
.getElementType()),
|
|
baseAddrAddr, /*members=*/mlir::SmallVector<mlir::Value>{},
|
|
/*member_index=*/mlir::DenseIntElementsAttr{}, op.getBounds(),
|
|
builder.getIntegerAttr(builder.getIntegerType(64, false),
|
|
op.getMapType().value()),
|
|
builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
|
|
mlir::omp::VariableCaptureKind::ByRef),
|
|
/*name=*/builder.getStringAttr(""),
|
|
/*partial_map=*/builder.getBoolAttr(false));
|
|
|
|
// TODO: map the addendum segment of the descriptor, similarly to the
|
|
// above base address/data pointer member.
|
|
|
|
if (auto mapClauseOwner =
|
|
llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
|
|
llvm::SmallVector<mlir::Value> newMapOps;
|
|
mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapOperands();
|
|
|
|
for (size_t i = 0; i < mapOperandsArr.size(); ++i) {
|
|
if (mapOperandsArr[i] == op) {
|
|
// Push new implicit maps generated for the descriptor.
|
|
newMapOps.push_back(baseAddr);
|
|
|
|
// for TargetOp's which have IsolatedFromAbove we must align the
|
|
// new additional map operand with an appropriate BlockArgument,
|
|
// as the printing and later processing currently requires a 1:1
|
|
// mapping of BlockArgs to MapInfoOp's at the same placement in
|
|
// each array (BlockArgs and MapOperands).
|
|
if (auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target))
|
|
targetOp.getRegion().insertArgument(i, baseAddr.getType(), loc);
|
|
}
|
|
newMapOps.push_back(mapOperandsArr[i]);
|
|
}
|
|
mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
|
|
}
|
|
|
|
mlir::Value newDescParentMapOp = builder.create<mlir::omp::MapInfoOp>(
|
|
op->getLoc(), op.getResult().getType(), descriptor,
|
|
mlir::TypeAttr::get(fir::unwrapRefType(descriptor.getType())),
|
|
/*varPtrPtr=*/mlir::Value{},
|
|
/*members=*/mlir::SmallVector<mlir::Value>{baseAddr},
|
|
/*members_index=*/
|
|
mlir::DenseIntElementsAttr::get(
|
|
mlir::VectorType::get(
|
|
llvm::ArrayRef<int64_t>({1, 1}),
|
|
mlir::IntegerType::get(builder.getContext(), 32)),
|
|
llvm::ArrayRef<int32_t>({0})),
|
|
/*bounds=*/mlir::SmallVector<mlir::Value>{},
|
|
builder.getIntegerAttr(builder.getIntegerType(64, false),
|
|
op.getMapType().value()),
|
|
op.getMapCaptureTypeAttr(), op.getNameAttr(), op.getPartialMapAttr());
|
|
op.replaceAllUsesWith(newDescParentMapOp);
|
|
op->erase();
|
|
}
|
|
|
|
// 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<mlir::omp::MapClauseOwningOpInterface>(target);
|
|
if (!mapClauseOwner)
|
|
return;
|
|
|
|
llvm::SmallVector<mlir::Value> newMapOps;
|
|
mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapOperands();
|
|
auto targetOp = llvm::dyn_cast<mlir::omp::TargetOp>(target);
|
|
|
|
for (size_t i = 0; i < mapOperandsArr.size(); ++i) {
|
|
if (mapOperandsArr[i] == op) {
|
|
for (auto [j, mapMember] : llvm::enumerate(op.getMembers())) {
|
|
newMapOps.push_back(mapMember);
|
|
// for TargetOp's which have IsolatedFromAbove we must align the
|
|
// new additional map operand with an appropriate BlockArgument,
|
|
// as the printing and later processing currently requires a 1:1
|
|
// mapping of BlockArgs to MapInfoOp's at the same placement in
|
|
// each array (BlockArgs and MapOperands).
|
|
if (targetOp) {
|
|
targetOp.getRegion().insertArgument(i + j, mapMember.getType(),
|
|
targetOp->getLoc());
|
|
}
|
|
}
|
|
}
|
|
newMapOps.push_back(mapOperandsArr[i]);
|
|
}
|
|
mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
|
|
}
|
|
|
|
// This pass executes on mlir::ModuleOp's finding 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.
|
|
void runOnOperation() override {
|
|
mlir::func::FuncOp func = getOperation();
|
|
mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
|
|
fir::KindMapping kindMap = fir::getKindMapping(module);
|
|
fir::FirOpBuilder builder{module, std::move(kindMap)};
|
|
|
|
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 per Target operation for the moment. However, when/if
|
|
// we optimise/cleanup the IR, it likely isn't too difficult to
|
|
// extend this function, it would require some modification to create a
|
|
// single new MapInfoOp per new MapInfoOp generated and share it across
|
|
// all users appropriately, making sure to only add a single member link
|
|
// per new generation for the original originating descriptor MapInfoOp.
|
|
assert(llvm::hasSingleElement(op->getUsers()) &&
|
|
"OMPMapInfoFinalization currently only supports single users "
|
|
"of a MapInfoOp");
|
|
|
|
if (!op.getMembers().empty()) {
|
|
addImplicitMembersToTarget(op, builder, *op->getUsers().begin());
|
|
} else if (fir::isTypeWithDescriptor(op.getVarType()) ||
|
|
mlir::isa_and_present<fir::BoxAddrOp>(
|
|
op.getVarPtr().getDefiningOp())) {
|
|
builder.setInsertionPoint(op);
|
|
genDescriptorMemberMaps(op, builder, *op->getUsers().begin());
|
|
}
|
|
});
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
namespace fir {
|
|
std::unique_ptr<mlir::Pass> createOMPMapInfoFinalizationPass() {
|
|
return std::make_unique<OMPMapInfoFinalizationPass>();
|
|
}
|
|
} // namespace fir
|