[llvm-branch-commits] [flang] [llvm] [MLIR][OpenMP] Add a new AutomapToTargetData conversion pass in FIR (PR #151989)

[Sergio Afonso via llvm-branch-commits]

================
@@ -0,0 +1,171 @@
+//===- AutomapToTargetData.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/Builder/DirectivesCommon.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/HLFIRTools.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include <mlir/Dialect/OpenMP/OpenMPInterfaces.h>
+#include <mlir/IR/Operation.h>
+
+namespace flangomp {
+#define GEN_PASS_DEF_AUTOMAPTOTARGETDATAPASS
+#include "flang/Optimizer/OpenMP/Passes.h.inc"
+} // namespace flangomp
+
+using namespace mlir;
+
+namespace {
+class AutomapToTargetDataPass
+    : public flangomp::impl::AutomapToTargetDataPassBase<
+          AutomapToTargetDataPass> {
+  // Returns true if the variable has a dynamic size and therefore requires
+  // bounds operations to describe its extents.
+  bool needsBoundsOps(Value var) {
+    assert(isa<omp::PointerLikeType>(var.getType()) &&
+           "only pointer like types expected");
+    Type t = fir::unwrapRefType(var.getType());
+    if (Type inner = fir::dyn_cast_ptrOrBoxEleTy(t))
+      return fir::hasDynamicSize(inner);
+    return fir::hasDynamicSize(t);
+  }
+
+  // Generate MapBoundsOp operations for the variable if required.
+  void genBoundsOps(fir::FirOpBuilder &builder, Value var,
+                    SmallVectorImpl<Value> &boundsOps) {
+    Location loc = var.getLoc();
+    fir::factory::AddrAndBoundsInfo info =
+        fir::factory::getDataOperandBaseAddr(builder, var,
+                                             /*isOptional=*/false, loc);
+    fir::ExtendedValue exv =
+        hlfir::translateToExtendedValue(loc, builder, hlfir::Entity{info.addr},
+                                        /*contiguousHint=*/true)
+            .first;
+    SmallVector<Value> tmp =
+        fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
+                                           mlir::omp::MapBoundsType>(
+            builder, info, exv, /*dataExvIsAssumedSize=*/false, loc);
+    llvm::append_range(boundsOps, tmp);
+  }
+
+  void findRelatedAllocmemFreemem(fir::AddrOfOp addressOfOp,
+                                  llvm::SmallVector<fir::StoreOp> &allocmems,
+                                  llvm::SmallVector<fir::LoadOp> &freemems) {
+    assert(addressOfOp->hasOneUse() && "op must have single use");
+
+    auto declaredRef =
+        cast<hlfir::DeclareOp>(*addressOfOp->getUsers().begin())->getResult(0);
+
+    for (Operation *refUser : declaredRef.getUsers()) {
+      if (auto storeOp = dyn_cast<fir::StoreOp>(refUser))
+        if (auto emboxOp = storeOp.getValue().getDefiningOp<fir::EmboxOp>())
+          if (auto allocmemOp =
+                  emboxOp.getOperand(0).getDefiningOp<fir::AllocMemOp>())
+            allocmems.push_back(storeOp);
+
+      if (auto loadOp = dyn_cast<fir::LoadOp>(refUser))
+        for (Operation *loadUser : loadOp.getResult().getUsers())
+          if (auto boxAddrOp = dyn_cast<fir::BoxAddrOp>(loadUser))
+            for (Operation *boxAddrUser : boxAddrOp.getResult().getUsers())
+              if (auto freememOp = dyn_cast<fir::FreeMemOp>(boxAddrUser))
+                freemems.push_back(loadOp);
+    }
+  }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation()->getParentOfType<ModuleOp>();
+    if (!module)
+      module = dyn_cast<ModuleOp>(getOperation());
+    if (!module)
+      return;
+
+    // Build FIR builder for helper utilities.
+    fir::KindMapping kindMap = fir::getKindMapping(module);
+    fir::FirOpBuilder builder{module, std::move(kindMap)};
+
+    // Collect global variables with AUTOMAP flag.
+    llvm::DenseSet<fir::GlobalOp> automapGlobals;
+    module.walk([&](fir::GlobalOp globalOp) {
+      if (auto iface =
+              dyn_cast<omp::DeclareTargetInterface>(globalOp.getOperation()))
+        if (iface.isDeclareTarget() && iface.getDeclareTargetAutomap())
+          automapGlobals.insert(globalOp);
+    });
+
+    for (fir::GlobalOp globalOp : automapGlobals)
+      if (auto uses = globalOp.getSymbolUses(module.getOperation()))
+        for (auto &x : *uses)
+          if (auto addrOp = dyn_cast<fir::AddrOfOp>(x.getUser())) {
+            llvm::SmallVector<fir::StoreOp> allocstores;
+            llvm::SmallVector<fir::LoadOp> freememloads;
+            findRelatedAllocmemFreemem(addrOp, allocstores, freememloads);
----------------
skatrak wrote:

Would it be possible to first gather all stores and loads for all uses and then process them? That way we wouldn't have to allocate/deallocate these lists for each use. Something like:
```c++
for (fir::GlobalOp globalOp : automapGlobals) {
  if (auto uses = globalOp.getSymbolUses(module.getOperation())) {
    llvm::SmallVector<fir::StoreOp> allocmemStores;
    llvm::SmallVector<fir::LoadOp> freememLoads;
    for (auto &x : *uses)
      if (auto addrOp = dyn_cast<fir::AddrOfOp>(x.getUser()))
        findRelatedAllocmemFreemem(addrOp, allocmemStores, freememLoads);
    for (auto storeOp : allocmemStores) { ... }
    for (auto loadOp : freememLoads) { ... }
  }
}
```

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