[llvm-branch-commits] [Flang][OpenMP][MLIR] Extend derived (record) type map support in Flang OpenMP by adding some initial support for explicit member mapping (PR #81511)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Wed Feb 14 09:06:38 PST 2024
================
@@ -0,0 +1,262 @@
+//===- 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 it's BlockArgument list
+/// to maintain consistency across all MapInfoOp's tied to a region directly or
+/// indirectly via an 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 = 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);
+
+ llvm::omp::OpenMPOffloadMappingFlags baseAddrMapFlag =
+ llvm::omp::OpenMPOffloadMappingFlags(op.getMapType().value());
+ baseAddrMapFlag |=
+ llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_PTR_AND_OBJ;
+
+ // 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, mlir::SmallVector<mlir::Value>{}, mlir::ArrayAttr{},
+ op.getBounds(),
+ builder.getIntegerAttr(
+ builder.getIntegerType(64, false),
+ static_cast<
+ std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+ baseAddrMapFlag)),
+ builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
+ mlir::omp::VariableCaptureKind::ByRef),
+ builder.getStringAttr("") /*name*/,
+ builder.getBoolAttr(false) /*partial_map*/);
+
+ // 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())),
+ mlir::Value{}, mlir::SmallVector<mlir::Value>{baseAddr},
+ mlir::ArrayAttr::get(builder.getContext(),
+ builder.getI64IntegerAttr(0)) /*members_index*/,
+ mlir::SmallVector<mlir::Value>{},
+ builder.getIntegerAttr(builder.getIntegerType(64, false),
+ op.getMapType().value()),
+ op.getMapCaptureTypeAttr(), op.getNameAttr(), op.getPartialMapAttr());
+ op.replaceAllUsesWith(newDescParentMapOp);
+ op->erase();
+ }
+
+ // For all mapped record members not directly used in the target region
+ // we add them to the block arguments in front of their parent and 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, only 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
+ // it helps to track 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 x and y together
+ // via the parent (at a kernel argument level in LLVM IR not just
+ // MLIR, 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, we do however, 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 unneccesary. It also doesn't
+ // quite make sense 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++ equivelant, which requires
+ // explicit mapping of these segments.
+ void addImplicitMembersToTarget(mlir::omp::MapInfoOp op,
+ fir::FirOpBuilder &builder,
+ mlir::Operation *target) {
+ 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 member maps
+ for (auto member : op.getMembers()) {
+ newMapOps.push_back(member);
+ // 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, member.getType(),
+ builder.getUnknownLoc());
+ }
+ }
+ newMapOps.push_back(mapOperandsArr[i]);
+ }
+ mapClauseOwner.getMapOperandsMutable().assign(newMapOps);
----------------
agozillon wrote:
It shouldn't at least from my understanding (which may be flawed), in the above case, we loop over the map operands of the owner (e.g. TargetOp) and generate a new vector of map operands to replace the previous, for every MapInfoOp in the map operands that is not the MapInfoOp we found in our walk and are currently processing we insert it into the new list at the same point, for the cases where it's the MapInfoOp we are processing we check if it has any members and directly insert them in-front of there parent on the list, and create a corresponding BlockArg if necessary, this new vector then replaces the old with the assign.
It's a little convoluted looking but unfortunately some of the Mutable interfaces you get from various MLIR types don't provide the best set of functionality, e.g. I don't believe in this particular case we can insert into the existing list at a particular index, which would be more concise, at least from what I've found, but perhaps I am missing some key thing!
I could add a test case for this in derived-type-map.f90 to double check all works as expected, there's a few runtime tests in one of the other related PRs that do this, but I haven't actually added the case when checking the IR itself.
This is very similar to what occurs for descriptors currently, so it's mimicked, however, with some hindsight it may have been better to walk the module for each of the map holding Target operations (TargetOp/EnterOp etc.), as opposed to the MapInfoOp, some of the code repetition required for it is likely alleviatable by templates, however, the main downside to that approach is that whenever we add a new map oeprand carrying operation . The upside is the pass would be a lot quicker, there's a lot of repetition with the current approach that's not always necessary. I'd leave a change like that for a follow up refactoring PR though personally as it touches the current descriptor mapping as well.
https://github.com/llvm/llvm-project/pull/81511
More information about the llvm-branch-commits
mailing list