[flang-commits] [flang] [flang][openacc] Add ACCHoistAllocatableRealloc pass (PR #206190)

[via flang-commits]

================
@@ -0,0 +1,208 @@
+//===- ACCHoistAllocatableRealloc.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
+//
+//===----------------------------------------------------------------------===//
+//
+// Hoist the allocatable reallocation that SeparateAllocatableAssign leaves
+// inside an OpenACC kernels construct out to the host, so only the assignment
+// runs on the device. A device-side allocate/deallocate is undefined behavior
+// under the OpenACC specification.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/OpenACC/Passes.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace fir::acc {
+#define GEN_PASS_DEF_ACCHOISTALLOCATABLEREALLOC
+#include "flang/Optimizer/OpenACC/Passes.h.inc"
+} // namespace fir::acc
+
+using namespace mlir;
+
+namespace {
+
+/// True for an allocatable descriptor reference:
+/// !fir.ref<!fir.box<!fir.heap<>>>.
+static bool isAllocatableDescRef(Type t) {
+  auto ref = dyn_cast<fir::ReferenceType>(t);
+  if (!ref)
+    return false;
+  auto box = dyn_cast<fir::BaseBoxType>(ref.getEleTy());
+  if (!box)
+    return false;
+  return isa<fir::HeapType>(box.getEleTy());
+}
+
+/// True if `v` is defined outside `region` (and therefore dominates the op that
+/// owns the region).
+static bool definedOutside(Value v, Region &region) {
+  return !region.isAncestor(v.getParentRegion());
+}
+
+/// Return the ancestor of `op` that is a direct child of `region`'s entry
+/// block, or null if `op` is not nested within that entry block.
+static Operation *topLevelInEntry(Operation *op, Region &region) {
+  Block *entry = &region.front();
+  for (Operation *cur = op; cur; cur = cur->getParentOp()) {
+    if (cur->getBlock() == entry)
+      return cur;
+  }
+  return nullptr;
+}
+
+/// Collect into `slice` the entry-block ops of `region` that `root` depends on
+/// (its backward SSA slice, including values used by nested ops). Returns false
+/// if a dependency is a region-local block argument, i.e. it cannot be hoisted.
+static bool collectReallocSlice(Operation *root, Region &region,
+                                llvm::SmallPtrSetImpl<Operation *> &slice) {
+  llvm::SmallVector<Operation *> worklist{root};
+  bool ok = true;
+  auto use = [&](Value v) {
+    if (definedOutside(v, region))
+      return;
+    if (Operation *def = v.getDefiningOp()) {
+      if (Operation *top = topLevelInEntry(def, region))
+        worklist.push_back(top);
+      else
+        ok = false;
+    } else {
+      ok = false; // region block argument
+    }
+  };
+  while (!worklist.empty()) {
+    Operation *op = worklist.pop_back_val();
+    if (!slice.insert(op).second)
+      continue;
+    for (Value v : op->getOperands())
+      use(v);
+    op->walk([&](Operation *nested) {
+      for (Value v : nested->getOperands())
+        use(v);
+    });
+  }
+  return ok;
+}
+
+/// True if hoisting the realloc (whose backward slice is `slice`, ending at
+/// `anchor`) past the kernel is safe. It is unsafe when a descriptor load that
+/// is not strictly after the realloc feeds an op that stays in the kernel: that
+/// op would read storage the realloc has already freed.
+static bool isSafeToHoist(Value desc, Operation *anchor, Region &region,
+                          const llvm::SmallPtrSetImpl<Operation *> &slice) {
+  for (Operation *user : desc.getUsers()) {
+    auto load = dyn_cast<fir::LoadOp>(user);
+    if (!load)
+      continue;
----------------
jeanPerier wrote:

I do not think we should ignore operations that are not fir.load. There could be a fir.convert/fir.declare hiding a fir.load....

I think the pass should bail if something is not a fir.load.

Also, this is making strong assumptions that the storage is not being aliased (that every access to the storage will be based on `desc` usage in the IR). This may not be true in case of TARGET allocatable whose data may be accessed via pointer, or module allocatables who can be access indirectly in function calls.
Maybe some of these cases can be ruled out inside kernel code, I am not sure about the rules there.

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