[flang-commits] [flang] [flang][acc] Add ACCDeclareActionConversion pass (PR #181894)

[via flang-commits]

================
@@ -0,0 +1,210 @@
+//===- ACCDeclareActionConversion.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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implements the allocation and deallocation semantics for allocatables and
+// pointers in declare directives. OpenACC 3.4, Section 2.13.2: in Fortran, if
+// a variable in the declare var-list has the allocatable or pointer attribute,
+// then for a non-shared memory device, an allocate (or intrinsic assignment
+// that allocates) allocates in both local and device memory and sets the
+// dynamic reference counter to one; a deallocate (or assignment that
+// deallocates) deallocates from both and sets the counter to zero.
+//
+// How this pass works:
+// - Lowering generates recipe functions that hold the recipe for creating the
+//   device copy (using acc dialect operations, e.g. acc.create).
+// - Lowering also attaches an attribute to the operations that allocate or
+//   deallocate the object.
+// - This pass finds operations with that attribute and inserts calls to the
+//   corresponding recipe.
+//
+// Example:
+//   module mm
+//     real, allocatable :: arr(:)
+//     !$acc declare create(arr)
+//   contains
+//     subroutine sub()
+//       allocate(arr(100))
+//     end subroutine sub
+//   end module mm
+//
+// Relevant IR before this pass (recipe function and store with attribute):
+//   func.func private @_QMmmEarr_acc_declare_update_desc_post_alloc(...) {
+//     ...  // acc ops to create/register device copy
+//     return
+//   }
+//   func.func @_QMmmPsub() {
+//     ...
+//     fir.store %box to %desc {acc.declare_action = #acc.declare_action<
+//       postAlloc = @_QMmmEarr_acc_declare_update_desc_post_alloc>} ...
+//   }
+//
+// After this pass (call to recipe inserted after the store):
+//   func.func @_QMmmPsub() {
+//     ...
+//     fir.store %box to %desc ...
+//     call @_QMmmEarr_acc_declare_update_desc_post_alloc()
+//   }
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/OpenACC/Passes.h"
+#include "flang/Optimizer/OpenACC/Support/FIROpenACCUtils.h"
+#include "flang/Runtime/entry-names.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "fir-acc-declare-action-conversion"
+
+namespace fir {
+namespace acc {
+#define GEN_PASS_DEF_ACCDECLAREACTIONCONVERSION
+#include "flang/Optimizer/OpenACC/Passes.h.inc"
+} // namespace acc
+} // namespace fir
+
+using namespace mlir;
+
+namespace {
+
+// Fortran runtime symbol names for pointer allocate/deallocate.
+static constexpr llvm::StringRef pointerAllocateName =
+    RTNAME_STRING(PointerAllocate);
+static constexpr llvm::StringRef pointerDeallocateName =
+    RTNAME_STRING(PointerDeallocate);
+
+class ACCDeclareActionConversion
+    : public fir::acc::impl::ACCDeclareActionConversionBase<
+          ACCDeclareActionConversion> {
+public:
+  using fir::acc::impl::ACCDeclareActionConversionBase<
+      ACCDeclareActionConversion>::ACCDeclareActionConversionBase;
+
+  void runOnOperation() override {
+    ModuleOp mod = getOperation();
+    OpBuilder builder(mod);
+
+    mod.walk([&](Operation *op) {
+      auto declareAction = op->getAttrOfType<acc::DeclareActionAttr>(
+          acc::getDeclareActionAttrName());
+      if (!declareAction)
+        return;
+
+      LLVM_DEBUG(llvm::dbgs() << "Found " << acc::getDeclareActionAttrName()
+                              << " on: " << *op << "\n");
+
+      auto preAlloc = declareAction.getPreAlloc();
+      auto postAlloc = declareAction.getPostAlloc();
+      auto preDealloc = declareAction.getPreDealloc();
+      auto postDealloc = declareAction.getPostDealloc();
+
+      if (!preAlloc && !postAlloc && !preDealloc && !postDealloc)
+        return;
+
+      for (auto action : {preAlloc, postAlloc, preDealloc, postDealloc}) {
+        if (!action)
+          continue;
+
+        if (auto func = dyn_cast<SymbolRefAttr>(action)) {
+          Operation *funcDef = SymbolTable::lookupNearestSymbolFrom(op, func);
----------------
jeanPerier wrote:

`SymbolTable::lookupNearestSymbolFrom` relies on a [linear visit](https://github.com/llvm/llvm-project/blob/051d125de4773951be5f2f51f245df94f7e29e2e/mlir/lib/IR/SymbolTable.cpp#L394) of the ModuleOp operation, so in big application, with many symbols, this may be very expensive to do it inside a visit (pseudo quadratic behavior).

I would advise moving the `lazySymbolTable` SymbolTable wrapper from https://github.com/llvm/llvm-project/blob/051d125de4773951be5f2f51f245df94f7e29e2e/flang/lib/Optimizer/Transforms/AbstractResult.cpp#L38 into a new flang/include/flang/Optimizer/Support/LazySymbolTable.h/.cpp and use that pattern here to only create the symbol table if this pass actually needs it (needs to do lookups), and get faster lookups (mlir::SymbolTable builds a map of the Symbol from the operation instead of going through it every time).

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