[Mlir-commits] [flang] [mlir] [flang][OpenMP] - Add `MapInfoOp` instances for target private variables when needed (PR #109862)

[llvmlistbot at llvm.org]

================
@@ -0,0 +1,156 @@
+//===- MapsForPrivatizedSymbols.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 creates MapInfoOp
+/// instances for certain privatized symbols.
+/// For example, if an allocatable variable is used in a private clause attached
+/// to a omp.target op, then the allocatable variable's descriptor will be
+/// needed on the device (e.g. GPU). This descriptor needs to be separately
+/// mapped onto the device. This pass creates the necessary omp.map.info ops for
+/// this.
+//===----------------------------------------------------------------------===//
+// TODO:
+// 1. Before adding omp.map.info, check if in case we already have an
+//    omp.map.info for the variable in question.
+// 2. Generalize this for more than just omp.target ops.
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.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/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include "llvm/Support/Debug.h"
+#include <type_traits>
+
+#define DEBUG_TYPE "omp-maps-for-privatized-symbols"
+
+namespace flangomp {
+#define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS
+#include "flang/Optimizer/OpenMP/Passes.h.inc"
+} // namespace flangomp
+using namespace mlir;
+namespace {
+class MapsForPrivatizedSymbolsPass
+    : public flangomp::impl::MapsForPrivatizedSymbolsPassBase<
+          MapsForPrivatizedSymbolsPass> {
+
+  bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) {
+    Region &allocRegion = privatizer.getAllocRegion();
+    Value blockArg0 = allocRegion.getArgument(0);
+    if (blockArg0.use_empty())
+      return false;
+    return true;
+  }
+  omp::MapInfoOp createMapInfo(Location loc, Value var,
+                               fir::FirOpBuilder &builder) {
+    uint64_t mapTypeTo = static_cast<
+        std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+        llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
+    Operation *definingOp = var.getDefiningOp();
+    auto declOp = llvm::dyn_cast_or_null<hlfir::DeclareOp>(definingOp);
+    assert(declOp &&
+           "Expected defining Op of privatized var to be hlfir.declare");
+
+    // We want the first result of the hlfir.declare op because our goal
+    // is to map the descriptor (fir.box or fir.boxchar) and the first
+    // result for hlfir.declare is the descriptor if a the symbol being
+    // decalred needs a descriptor.
+    Value varPtr = declOp.getBase();
+
+    // If we do not have a reference to descritor, but the descriptor itself
+    // then we need to store that on the stack so that we can map the
+    // address of the descriptor.
+    if (mlir::isa<fir::BaseBoxType>(varPtr.getType()) ||
+        mlir::isa<fir::BoxCharType>(varPtr.getType())) {
+      OpBuilder::InsertPoint savedInsPoint = builder.saveInsertionPoint();
+      mlir::Block *allocaBlock = builder.getAllocaBlock();
+      assert(allocaBlock && "No allocablock  found for a funcOp");
+      builder.setInsertionPointToStart(allocaBlock);
+      auto alloca = builder.create<fir::AllocaOp>(loc, varPtr.getType());
----------------
agozillon wrote:

It may be important in the future to keep track of the allocas in scope made for descriptors and re-use them where there's an alloca in scope for that variable that we've already assigned and stored too, as it lead to a couple of issues in larger examples for at least the map case, I don't know if it'd be an issue here, so it's just something to keep in mind incase you run into any issues rather than a change suggestions :-)

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