[Mlir-commits] [flang] [mlir] [flang][acc] Add ACCOptimizeFirstprivateMap pass (PR #178546)

[llvmlistbot at llvm.org]

================
@@ -0,0 +1,190 @@
+//===- ACCOptimizeFirstprivateMap.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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass optimizes firstprivate mapping operations (acc.firstprivate_map).
+// The optimization hoists loads from the firstprivate variable to before the
+// compute region, effectively converting the firstprivate copy to a
+// pass-by-value pattern. This eliminates the need for runtime copying into
+// global memory.
+//
+// Example transformation:
+//
+//   Before:
+//     %decl = fir.declare %alloca : !fir.ref<i32>
+//     %fp = acc.firstprivate_map varPtr(%decl) -> !fir.ref<i32>
+//     acc.parallel {
+//       %val = fir.load %fp : !fir.ref<i32>  // load inside region
+//       ...
+//     }
+//
+//   After:
+//     %decl = fir.declare %alloca : !fir.ref<i32>
+//     %val = fir.load %decl : !fir.ref<i32>  // load hoisted before region
+//     acc.parallel {
+//       ...  // uses %val directly
+//     }
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/FortranVariableInterface.h"
+#include "flang/Optimizer/OpenACC/Passes.h"
+#include "flang/Optimizer/OpenACC/Support/FIROpenACCUtils.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace fir::acc {
+#define GEN_PASS_DEF_ACCOPTIMIZEFIRSTPRIVATEMAP
+#include "flang/Optimizer/OpenACC/Passes.h.inc"
+} // namespace fir::acc
+
+using namespace mlir;
+
+namespace {
+
+/// Returns the enclosing offload region interface, or nullptr if not inside
+/// one.
+static acc::OffloadRegionOpInterface getEnclosingOffloadRegion(Operation *op) {
+  Operation *parent = op->getParentOp();
+  while (parent) {
+    if (auto offloadOp = dyn_cast<acc::OffloadRegionOpInterface>(parent))
+      return offloadOp;
+    parent = parent->getParentOp();
+  }
+  return nullptr;
+}
+
+/// Returns true if the value is defined by an OpenACC data clause operation.
+static bool isDefinedByDataClause(Value value) {
+  Operation *defOp = value.getDefiningOp();
+  if (!defOp)
+    return false;
+  return acc::getDataClause(defOp).has_value();
+}
+
+/// Returns true if the value is defined inside the given offload region.
+/// This handles both operation results and block arguments.
+static bool isDefinedInsideRegion(Value value,
+                                  acc::OffloadRegionOpInterface offloadOp) {
+  Region *valueRegion = value.getParentRegion();
+  if (!valueRegion)
+    return false;
+  return offloadOp.getOffloadRegion().isAncestor(valueRegion);
+}
+
+/// Returns true if the variable may be optional.
+static bool mayBeOptionalVariable(Value var) {
+  // Don't strip declare ops - we need to check the optional attribute on them.
+  Value originalDef = fir::acc::getOriginalDef(var, /*stripDeclare=*/false);
+  if (auto varIface = dyn_cast_or_null<fir::FortranVariableOpInterface>(
+          originalDef.getDefiningOp()))
+    return varIface.isOptional();
+  // If the defining op is an alloca, it's a local variable and not optional.
+  if (isa_and_nonnull<fir::AllocaOp, fir::AllocMemOp>(
+          originalDef.getDefiningOp()))
+    return false;
+  // Conservative: if we can't determine, assume it may be optional.
+  return true;
+}
+
+/// Returns true if the type is a reference to a trivial type.
+/// Note that this does not allow fir.heap, fir.ptr, or fir.llvm_ptr
+/// types - since we would need to check if the load is valid via
+/// a null-check to enable the optimization.
+static bool isRefToTrivialType(Type type) {
+  if (!mlir::isa<fir::ReferenceType>(type))
+    return false;
+  return fir::isa_trivial(fir::unwrapRefType(type));
+}
+
+static void hoistLoads(acc::FirstprivateMapInitialOp firstprivateInitOp,
+                       Value var, Value accVar) {
+  llvm::SmallVector<fir::LoadOp> loadsToHoist;
+  for (Operation *user : accVar.getUsers()) {
+    if (auto loadOp = dyn_cast<fir::LoadOp>(user))
+      loadsToHoist.push_back(loadOp);
----------------
jeanPerier wrote:

It seems the pass is assuming there are only loads made to the accVar (otherwise it would be invalid to delete the firstprivateInitOp).

I think in practice that is what happening given how FirstprivateMapInitialOp are generated, but it would be safer to bail here if any user is not a load.

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