95fe47ca7e
This patch seeks to add an initial lowering for pointers and allocatable variables captured by implicit and explicit map in Flang OpenMP for Target operations that take map clauses e.g. Target, Target Update. Target Exit/Enter etc. Currently this is done by treating the type that lowers to a descriptor (allocatable/pointer/assumed shape) as a map of a record type (e.g. a structure) as that's effectively what descriptor types lower to in LLVM-IR and what they're represented as in the Fortran runtime (written in C/C++). The descriptor effectively lowers to a structure containing scalar and array elements that represent various aspects of the underlying data being mapped (lower bound, upper bound, extent being the main ones of interest in most cases) and a pointer to the allocated data. In this current iteration of the mapping we map the structure in it's entirety and then attach the underlying data pointer and map the data to the device, this allows most of the required data to be resident on the device for use. Currently we do not support the addendum (another block of pointer data), but it shouldn't be too difficult to extend this to support it. The MapInfoOp generation for descriptor types is primarily handled in an optimization pass, where it expands BoxType (descriptor types) map captures into two maps, one for the structure (scalar elements) and the other for the pointer data (base address) and links them in a Parent <-> Child relationship. The later lowering processes will then treat them as a conjoined structure with a pointer member map.
169 lines
7.5 KiB
C++
169 lines
7.5 KiB
C++
//===- OMPDescriptorMapInfoGen.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 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.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#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 <iterator>
|
|
|
|
namespace fir {
|
|
#define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS
|
|
#include "flang/Optimizer/Transforms/Passes.h.inc"
|
|
} // namespace fir
|
|
|
|
namespace {
|
|
class OMPDescriptorMapInfoGenPass
|
|
: public fir::impl::OMPDescriptorMapInfoGenPassBase<
|
|
OMPDescriptorMapInfoGenPass> {
|
|
|
|
void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
|
|
fir::FirOpBuilder &builder,
|
|
mlir::Operation *target) {
|
|
mlir::Location loc = builder.getUnknownLoc();
|
|
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,
|
|
llvm::cast<mlir::omp::PointerLikeType>(
|
|
fir::unwrapRefType(baseAddrAddr.getType()))
|
|
.getElementType(),
|
|
baseAddrAddr, mlir::SmallVector<mlir::Value>{}, op.getBounds(),
|
|
builder.getIntegerAttr(builder.getIntegerType(64, false),
|
|
op.getMapType().value()),
|
|
builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
|
|
mlir::omp::VariableCaptureKind::ByRef),
|
|
builder.getStringAttr("") /*name*/);
|
|
|
|
// 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,
|
|
fir::unwrapRefType(descriptor.getType()), mlir::Value{},
|
|
mlir::SmallVector<mlir::Value>{baseAddr},
|
|
mlir::SmallVector<mlir::Value>{},
|
|
builder.getIntegerAttr(builder.getIntegerType(64, false),
|
|
op.getMapType().value()),
|
|
op.getMapCaptureTypeAttr(), op.getNameAttr());
|
|
op.replaceAllUsesWith(newDescParentMapOp);
|
|
op->erase();
|
|
}
|
|
|
|
// 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) {
|
|
if (fir::isTypeWithDescriptor(op.getVarType()) ||
|
|
mlir::isa_and_present<fir::BoxAddrOp>(
|
|
op.getVarPtr().getDefiningOp())) {
|
|
builder.setInsertionPoint(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()) &&
|
|
"OMPDescriptorMapInfoGen currently only supports single users "
|
|
"of a MapInfoOp");
|
|
genDescriptorMemberMaps(op, builder, *op->getUsers().begin());
|
|
}
|
|
});
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
namespace fir {
|
|
std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass() {
|
|
return std::make_unique<OMPDescriptorMapInfoGenPass>();
|
|
}
|
|
} // namespace fir
|