[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)

[ioana ghiban]

================
@@ -0,0 +1,126 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/LogicalResult.h"
+
+namespace mlir {
+namespace memref {
+#define GEN_PASS_DEF_ELEVATEALLOCSTOGLOBALSPASS
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+} // namespace memref
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+/// Returns true if `op` is contained inside any branching, region, or looping
+/// structure (such as scf.for, scf.if, or repetitive regions)
+static bool isInsideControlFlow(Operation *op) {
+  return getEnclosingRepetitiveRegion(op) ||
+         op->getParentOfType<LoopLikeOpInterface>() ||
+         op->getParentOfType<RegionBranchOpInterface>();
+}
+
+/// Elevates a static `memref.alloc` operation to a top-level `memref.global` op
+/// if the allocation is not enclosed within any control flow constructs.
+///
+/// Converts:
+/// ```mlir
+/// %0 = memref.alloc() : memref<4x4xf32>
+/// memref.dealloc %0 : memref<4x4xf32>
+/// ```
+/// to:
+/// ```mlir
+/// memref.global "private" @global_alloc : memref<4x4xf32>
+/// ...
+/// %0 = memref.get_global @global_alloc : memref<4x4xf32>
+/// ```
+struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::AllocOp allocOp,
+                                PatternRewriter &rewriter) const final {
+    auto memrefType = allocOp.getType();
+    // `memref.global` requires statically shaped memrefs with no dynamic sizes.
+    if (!memrefType.hasStaticShape() || !allocOp.getDynamicSizes().empty())
+      return failure();
+
+    // Avoid elevating allocations inside control flow (loops or conditionals),
+    // as converting them to a single static global would make multiple
+    // executions share the same buffer, changing semantics or causing race
+    // conditions.
+    if (isInsideControlFlow(allocOp))
----------------
ioghiban wrote:

This transformation changes a fresh allocation into one persistent storage instance. The current eligibility checks only consider some intraprocedural control-flow nesting, but they do not establish that the containing function is invoked only once.

For example, two sequential calls to the same function would originally receive distinct allocations, while after this transformation both calls use the same global. The same issue applies to recursion and concurrent invocation.

What pipeline invariant makes this storage-class and allocation-identity change valid, for the considered scope? If the bufferized model entry point is guaranteed to execute once and non-concurrently, could that precondition be documented or represented explicitly? Otherwise, the pass seems to need interprocedural eligibility checks in addition to the local control-flow checks.

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