[flang-commits] [flang] [mlir] [Flang][mlir] - Translation of delayed privatization for deferred target-tasks (PR #155348)

[Pranav Bhandarkar via flang-commits]

================
@@ -0,0 +1,423 @@
+//===- OpenMPOffloadPrivatizationPrepare.cpp - Prepare for OpenMP Offload
+// Privatization ---------===//
+//
+// 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 "mlir/Dialect/LLVMIR/Transforms/OpenMPOffloadPrivatizationPrepare.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include <cstdint>
+#include <utility>
+
+//===----------------------------------------------------------------------===//
+// A pass that prepares OpenMP code for translation of delayed privatization
+// in the context of deferred target tasks. Deferred target tasks are created
+// when the nowait clause is used on the target directive.
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "omp-prepare-for-offload-privatization"
+#define PDBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "]: ")
+
+namespace mlir {
+namespace LLVM {
+
+#define GEN_PASS_DEF_PREPAREFOROMPOFFLOADPRIVATIZATIONPASS
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
+
+} // namespace LLVM
+} // namespace mlir
+
+using namespace mlir;
+namespace {
+
+//===----------------------------------------------------------------------===//
+// OMPTargetPrepareDelayedPrivatizationPattern
+//===----------------------------------------------------------------------===//
+
+class OMPTargetPrepareDelayedPrivatizationPattern
+    : public OpRewritePattern<omp::TargetOp> {
+public:
+  using OpRewritePattern<omp::TargetOp>::OpRewritePattern;
+
+  // Match omp::TargetOp that have the following characteristics.
+  // 1. have private vars which refer to local (stack) memory
+  // 2. the target op has the nowait clause
+  // In this case, we allocate memory for the privatized variable on the heap
+  // and copy the original variable into this new heap allocation. We fix up
+  // any omp::MapInfoOp instances that may be mapping the private variable.
+  mlir::LogicalResult
+  matchAndRewrite(omp::TargetOp targetOp,
+                  PatternRewriter &rewriter) const override {
+    if (!hasPrivateVars(targetOp) || !isTargetTaskDeferred(targetOp))
+      return rewriter.notifyMatchFailure(
+          targetOp,
+          "targetOp does not have privateVars or does not need a target task");
+
+    ModuleOp mod = targetOp->getParentOfType<ModuleOp>();
+    LLVM::LLVMFuncOp llvmFunc = targetOp->getParentOfType<LLVM::LLVMFuncOp>();
+    OperandRange privateVars = targetOp.getPrivateVars();
+    mlir::SmallVector<mlir::Value> newPrivVars;
+
+    newPrivVars.reserve(privateVars.size());
+    std::optional<ArrayAttr> privateSyms = targetOp.getPrivateSyms();
+    for (auto [privVarIdx, privVarSymPair] :
+         llvm::enumerate(llvm::zip_equal(privateVars, *privateSyms))) {
+      auto privVar = std::get<0>(privVarSymPair);
+      auto privSym = std::get<1>(privVarSymPair);
+
+      omp::PrivateClauseOp privatizer = findPrivatizer(targetOp, privSym);
+      if (!privatizer.needsMap()) {
+        newPrivVars.push_back(privVar);
+        continue;
+      }
+      bool isFirstPrivate = privatizer.getDataSharingType() ==
+                            omp::DataSharingClauseType::FirstPrivate;
+
+      mlir::Value mappedValue =
+          targetOp.getMappedValueForPrivateVar(privVarIdx);
+      Operation *mapInfoOperation = mappedValue.getDefiningOp();
+      auto mapInfoOp = mlir::cast<omp::MapInfoOp>(mapInfoOperation);
+
+      if (mapInfoOp.getMapCaptureType() == omp::VariableCaptureKind::ByCopy) {
+        newPrivVars.push_back(privVar);
+        continue;
+      }
+
+      // Allocate heap memory that corresponds to the type of memory
+      // pointed to by varPtr
+      // TODO: For boxchars this likely wont be a pointer.
+      mlir::Value varPtr = privVar;
+      mlir::Value heapMem = allocateHeapMem(targetOp, privVar, mod, rewriter);
+      if (!heapMem)
+        return failure();
+
+      newPrivVars.push_back(heapMem);
+
+      // Find the earliest insertion point for the copy. This will be before
+      // the first in the list of omp::MapInfoOp instances that use varPtr.
+      // After the copy these omp::MapInfoOp instances will refer to heapMem
+      // instead.
+      Operation *varPtrDefiningOp = varPtr.getDefiningOp();
+      std::set<Operation *> users;
+      users.insert(varPtrDefiningOp->user_begin(),
+                   varPtrDefiningOp->user_end());
+
+      auto usesVarPtr = [&users](Operation *op) -> bool {
+        return users.count(op);
+      };
+      SmallVector<Operation *> chainOfOps;
+      chainOfOps.push_back(mapInfoOperation);
+      if (!mapInfoOp.getMembers().empty()) {
+        for (auto member : mapInfoOp.getMembers()) {
+          if (usesVarPtr(member.getDefiningOp()))
+            chainOfOps.push_back(member.getDefiningOp());
+
+          omp::MapInfoOp memberMap =
+              mlir::cast<omp::MapInfoOp>(member.getDefiningOp());
+          if (memberMap.getVarPtrPtr() &&
+              usesVarPtr(memberMap.getVarPtrPtr().getDefiningOp()))
+            chainOfOps.push_back(memberMap.getVarPtrPtr().getDefiningOp());
+        }
+      }
+      DominanceInfo dom;
+      llvm::sort(chainOfOps, [&](Operation *l, Operation *r) {
+        return dom.dominates(l, r);
+      });
+
+      rewriter.setInsertionPoint(chainOfOps.front());
+      // Copy the value of the local variable into the heap-allocated location.
+      mlir::Location loc = chainOfOps.front()->getLoc();
+      mlir::Type varType = getElemType(varPtr);
+      auto loadVal = rewriter.create<LLVM::LoadOp>(loc, varType, varPtr);
----------------
bhandarkar-pranav wrote:

You are right that there is a problem here. The problem is that derived types could have any number of pointers and therefore deep copies will be needed. Are you suggesting that I use the copy region of the privatizer by cloning it to get the deep copy that i need? Before that though, I'd have to allocate memory for each pointer inside the derived type. I was hoping to tackle derived types in a subsequent PR, which I should have made clear in this PR.

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