[llvm-branch-commits] [flang] [flang] Add policy-driven allocation-placement pass - memory passes unification [2/5] (PR #210742)

[Tom Eccles via llvm-branch-commits]

================
@@ -0,0 +1,324 @@
+//===- AllocationPlacement.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 decides, for each array allocation in a function, whether it should
+// live on the stack (fir.alloca) or on the heap (fir.allocmem), and rewrites it
+// accordingly. The decision is delegated to the policy in
+// AllocationPlacementPolicy.h. Two rewrite engines are reused:
+//   - stack-to-heap uses fir::replaceAllocas (MemoryUtils);
+//   - heap-to-stack reuses the StackArrays analysis and rewrite, which only
+//     stackifies fir.allocmem that are provably freed on all paths.
+//
+//===----------------------------------------------------------------------===//
+
+#include "StackArrays.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "flang/Optimizer/Support/DataLayout.h"
+#include "flang/Optimizer/Transforms/AllocationPlacementPolicy.h"
+#include "flang/Optimizer/Transforms/MemoryUtils.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+#include <optional>
+
+namespace fir {
+#define GEN_PASS_DEF_ALLOCATIONPLACEMENT
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "allocation-placement"
+
+//===----------------------------------------------------------------------===//
+// Default placement policy
+//===----------------------------------------------------------------------===//
+
+fir::AllocationPlacement
+fir::decideAllocationPlacement(const AllocationInfo &info,
+                               const AllocationPlacementThresholds &thresholds,
+                               std::size_t stackBytesUsed) {
+  using P = fir::AllocationPlacement;
+
+  // Translate a "should this be on the stack" decision into a placement,
+  // accounting for where the allocation currently lives.
+  auto place = [&](bool wantStack) -> P {
+    if (wantStack)
+      return info.isCurrentlyOnStack ? P::Leave : P::Stack;
+    return info.isCurrentlyOnStack ? P::Heap : P::Leave;
+  };
+
+  // -fstack-arrays: put everything on the stack (best effort). The
+  // heap-to-stack conversion still only happens where it is provably safe.
+  if (thresholds.stackArrays)
+    return place(/*wantStack=*/true);
+
+  // Runtime-sized arrays (automatic arrays, dynamic temporaries) go on the
+  // heap.
+  if (info.isDynamic)
+    return place(/*wantStack=*/false);
+
+  // Without a known constant size we cannot reason about thresholds.
+  if (!info.byteSize)
+    return P::Leave;
+
+  // Constant-size user variables always go on the stack.
+  if (!info.isTemporary)
+    return place(/*wantStack=*/true);
+
+  auto size = static_cast<std::size_t>(*info.byteSize);
+  if (size <= thresholds.smallArrayThresholdBytes)
+    // Small arrays go on the stack while the per-function budget allows it.
+    return place(/*wantStack=*/stackBytesUsed + size <=
+                 thresholds.totalStackLimitBytes);
+
+  // Big array temporaries go on the heap.
+  return place(/*wantStack=*/false);
+}
+
+namespace {
+
+/// Return true if the allocation is a compiler temporary, i.e. it has no
+/// uniqued name (user variables always carry one).
+static bool isTemporaryAllocation(mlir::Operation *op) {
----------------
tblah wrote:

what about checking for `fir.must_be_heap` on allocmem operations?

Or if we don't need this at all perhaps we should delete that attribute.

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