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

[Tom Eccles 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()) {
----------------
tblah wrote:

For me this would be easier to understand using privatizer.readsFromMold(), but this might be obvious to someone who works on target offloading regularly so feel free to leave it as it is if you prefer.

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