[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)

Kareem Ergawy via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Thu Feb 15 06:38:00 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);
----------------
ergawy wrote:

Thanks for the clarification, makes sense.

https://github.com/llvm/llvm-project/pull/81511


More information about the llvm-branch-commits mailing list