Add MapsForPrivatizedSymbolsPass

Author: bhandarkar-pranav

flang/include/flang/Optimizer/OpenMP/Passes.td

@@ -22,6 +22,19 @@ def MapInfoFinalizationPass
   let dependentDialects = ["mlir::omp::OpenMPDialect"];
 }
 
+def MapsForPrivatizedSymbolsPass
+    : Pass<"omp-map-for-privatized-symbols", "mlir::func::FuncOp"> {
+  let summary = "Creates MapInfoOp instances for privatized symbols when needed";
+  let description = [{
+    Adds omp.map.info operations for privatized symbols on omp.target ops
+    In certain situations, such as when an allocatable is privatized, its
+    descriptor is needed in the alloc region of the privatizer. This results
+    in the use of the descriptor inside the target region. As such, the
+    descriptor then needs to be mapped. This pass adds such MapInfoOp operations.
+  }];
+  let dependentDialects = ["mlir::omp::OpenMPDialect"];
+}
+
 def MarkDeclareTargetPass
     : Pass<"omp-mark-declare-target", "mlir::ModuleOp"> {
   let summary = "Marks all functions called by an OpenMP declare target function as declare target";

flang/include/flang/Tools/CLOptions.inc

@@ -368,6 +368,7 @@ inline void createHLFIRToFIRPassPipeline(
 inline void createOpenMPFIRPassPipeline(
     mlir::PassManager &pm, bool isTargetDevice) {
   pm.addPass(flangomp::createMapInfoFinalizationPass());
+  pm.addPass(flangomp::createMapsForPrivatizedSymbolsPass());
   pm.addPass(flangomp::createMarkDeclareTargetPass());
   if (isTargetDevice)
     pm.addPass(flangomp::createFunctionFilteringPass());

flang/lib/Optimizer/OpenMP/CMakeLists.txt

@@ -2,6 +2,7 @@ get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 add_flang_library(FlangOpenMPTransforms
   FunctionFiltering.cpp
+  MapsForPrivatizedSymbols.cpp
   MapInfoFinalization.cpp
   MarkDeclareTarget.cpp
 

flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp

@@ -0,0 +1,119 @@
+//===- 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
+//
+//===----------------------------------------------------------------------===//
+#include "flang/Optimizer/Dialect/FIRType.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 <type_traits>
+
+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;
+  }
+  void dumpPrivatizerInfo(omp::PrivateClauseOp &privatizer,
+                          mlir::Value privVar) {
+    llvm::errs() << "Found a privatizer:\n";
+    privatizer.dump();
+    llvm::errs() << "\n";
+
+    llvm::errs() << "$type = " << privatizer.getType() << "\n";
+    llvm::errs() << "privVar = ";
+    privVar.dump();
+    llvm::errs() << "\n";
+
+    llvm::errs() << "privVar.getDefiningOp() = ";
+    privVar.getDefiningOp()->dump();
+    llvm::errs() << "\n";
+    llvm::errs() << "\n";
+  }
+  omp::MapInfoOp createMapInfo(mlir::Location loc, mlir::Value var,
+                               OpBuilder &builder) {
+    //    llvm::omp::OpenMPOffloadMappingFlags mapTypeBits =
+    //    llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
+    uint64_t mapTypeTo = static_cast<
+        std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+        llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
+    return builder.create<omp::MapInfoOp>(
+        loc, var.getType(), var,
+        mlir::TypeAttr::get(fir::unwrapRefType(var.getType())),
+        /*varPtrPtr=*/mlir::Value{},
+        /*members=*/mlir::SmallVector<mlir::Value>{},
+        /*member_index=*/mlir::DenseIntElementsAttr{},
+        /*bounds=*/mlir::ValueRange{},
+        builder.getIntegerAttr(builder.getIntegerType(64, /*isSigned=*/false),
+                               mapTypeTo),
+        builder.getAttr<omp::VariableCaptureKindAttr>(
+            omp::VariableCaptureKind::ByRef),
+        mlir::StringAttr(), builder.getBoolAttr(false));
+  }
+  void addMapInfoOp(omp::TargetOp targetOp, omp::MapInfoOp mapInfoOp) {
+    mlir::Location loc = targetOp.getLoc();
+    targetOp.getMapVarsMutable().append(mlir::ValueRange{mapInfoOp});
+    size_t numMapVars = targetOp.getMapVars().size();
+    targetOp.getRegion().insertArgument(numMapVars - 1, mapInfoOp.getType(),
+                                        loc);
+  }
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    OpBuilder builder(context);
+    getOperation()->walk([&](omp::TargetOp targetOp) {
+      llvm::errs() << "MapsForPrivatizedSymbolsPass::TargetOp is \n";
+      targetOp.dump();
+      llvm::errs() << "\n";
+
+      if (targetOp.getPrivateVars().empty())
+        return;
+
+      OperandRange privVars = targetOp.getPrivateVars();
+      std::optional<ArrayAttr> privSyms = targetOp.getPrivateSyms();
+
+      for (auto [privVar, privSym] : llvm::zip_equal(privVars, *privSyms)) {
+
+        SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym);
+        omp::PrivateClauseOp privatizer =
+            SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(
+                targetOp, privatizerName);
+
+        assert(mlir::isa<fir::ReferenceType>(privVar.getType()) &&
+               "Privatized variable should be a reference.");
+        if (!privatizerNeedsMap(privatizer)) {
+          return;
+        }
+        llvm::errs() << "Privatizer NEEDS a map\n";
+        builder.setInsertionPoint(targetOp);
+        dumpPrivatizerInfo(privatizer, privVar);
+
+        mlir::Location loc = targetOp.getLoc();
+        omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder);
+        addMapInfoOp(targetOp, mapInfoOp);
+        llvm::errs() << __FUNCTION__ << "MapInfoOp is \n";
+        mapInfoOp.dump();
+        llvm::errs() << "\n";
+      }
+    });
+  }
+};
+} // namespace

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -880,7 +880,7 @@ def MapInfoOp : OpenMP_Op<"map.info", [AttrSizedOperandSegments]> {
        objects (e.g. derived types or classes), indicates the bounds to be copied
        of the variable. When it's an array slice it is in rank order where rank 0
        is the inner-most dimension.
-    - 'map_clauses': OpenMP map type for this map capture, for example: from, to and
+    - 'map_type': OpenMP map type for this map capture, for example: from, to and
        always. It's a bitfield composed of the OpenMP runtime flags stored in
        OpenMPOffloadMappingFlags.
     - 'map_capture_type': Capture type for the variable e.g. this, byref, byvalue, byvla