[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 23 13:54:59 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Bhavesh M (beamandala)
<details>
<summary>Changes</summary>
This creates the `elevate-allocs-to-globals` pass which replaces `memref::AllocOp`s with `memref::GlobalOp`s in the closest symbol table defining op if it's a `ModuleOp`. It does this through the following process:
- Check that the `AllocOp` is statically shaped because `GlobalOp` requires static shape.
- Check that the `AllocOp` isn't in control flow such as loops or conditionals because elevating in these scenarios could change program semantics.
- Check that the closest symbol table defining op is a `ModuleOp` because other `GlobalOp`s could be illegal in other symbol table defining ops and to minimize complexity for initial implementation.
- Create the `GlobalOp`, delete any `memref::DeallocOp`s associated with the `AllocOp`
- Replace the `AllocOp` with `memref::GetGlobalOp`
This pass will help reduce/eliminate dynamic memory allocations in frequently executed functions.
---
Full diff: https://github.com/llvm/llvm-project/pull/211141.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td (+34)
- (modified) mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h (+3)
- (modified) mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt (+1)
- (added) mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp (+126)
- (added) mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir (+140)
``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index c6be600247696..a0a4be7318889 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -340,4 +340,38 @@ def FlattenMemrefsPass : Pass<"flatten-memref"> {
];
}
+def ElevateAllocsToGlobalsPass : Pass<"elevate-allocs-to-globals", "ModuleOp"> {
+ let summary = "Elevate allocs to globals";
+ let description = [{
+ This pass converts statically-shaped `memref.alloc` operations that are not
+ enclosed within loops or control flow into private `memref.global` operations.
+ `memref.alloc` operations are replaced with `memref.get_global` references and
+ any associated `memref.dealloc` operations are removed.
+
+ Example:
+
+ ```mlir
+ func.func @example() {
+ %0 = memref.alloc() {alignment = 64 : i64} : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
+ return
+ }
+ ```
+
+ is transformed to
+
+ ```mlir
+ memref.global "private" @global_alloc : memref<10xf32> = uninitialized {alignment = 64 : i64}
+
+ func.func @example() {
+ %0 = memref.get_global @global_alloc : memref<10xf32>
+ return
+ }
+ ```
+ }];
+ let dependentDialects = [
+ "memref::MemRefDialect"
+ ];
+}
+
#endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
index 720677455ae5d..6960b8ef00f06 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
@@ -14,6 +14,7 @@
#ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
#define MLIR_DIALECT_MEMREF_TRANSFORMS_TRANSFORMS_H
+#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLFunctionalExtras.h"
@@ -165,6 +166,8 @@ void populateExtractAddressComputationsPatterns(RewritePatternSet &patterns);
/// into one-dimensional memref operations.
void populateFlattenMemrefsPatterns(RewritePatternSet &patterns);
+void populateElevateAllocsToGlobalsPatterns(RewritePatternSet &patterns);
+
/// Build a new memref::AllocaOp whose dynamic sizes are independent of all
/// given independencies. If the op is already independent of all
/// independencies, the same AllocaOp result is returned.
diff --git a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
index 1a8b03dabcfb7..7d43e66989bc4 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
@@ -17,6 +17,7 @@ add_mlir_dialect_library(MLIRMemRefTransforms
ReifyResultShapes.cpp
ResolveShapedTypeResultDims.cpp
RuntimeOpVerification.cpp
+ ElevateAllocsToGlobals.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
new file mode 100644
index 0000000000000..907038b545c62
--- /dev/null
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -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))
+ return failure();
+
+ // Create the global variable at the nearest enclosing symbol table defining
+ // op if it's a ModuleOp.
+ auto moduleOp = llvm::dyn_cast_or_null<ModuleOp>(
+ SymbolTable::getNearestSymbolTable(allocOp));
+ if (!moduleOp)
+ return failure();
+
+ OpBuilder detachedBuilder(rewriter.getContext());
+ StringAttr globalName = rewriter.getStringAttr("global_alloc");
+ memref::GlobalOp globalOp = memref::GlobalOp::create(
+ detachedBuilder, allocOp.getLoc(), globalName,
+ rewriter.getStringAttr("private"), memrefType, rewriter.getUnitAttr(),
+ false, allocOp.getAlignmentAttr());
+
+ SymbolTable(moduleOp).insert(globalOp);
+
+ // Remove any `memref.dealloc` operations using this allocation
+ SmallVector<Operation *> deallocsToDelete;
+ for (OpOperand &use : allocOp.getResult().getUses()) {
+ Operation *user = use.getOwner();
+ if (isa<memref::DeallocOp>(user))
+ deallocsToDelete.push_back(user);
+ }
+ for (Operation *dealloc : deallocsToDelete)
+ rewriter.eraseOp(dealloc);
+
+ // Replace the original `memref.alloc` with `memref.get_global`.
+ rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(allocOp, memrefType,
+ globalOp.getName());
+
+ return success();
+ }
+};
+
+struct ElevateAllocsToGlobalsPass
+ : public mlir::memref::impl::ElevateAllocsToGlobalsPassBase<
+ ElevateAllocsToGlobalsPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ ModuleOp moduleOp = getOperation();
+
+ RewritePatternSet patterns(&getContext());
+ memref::populateElevateAllocsToGlobalsPatterns(patterns);
+
+ (void)applyPatternsGreedily(moduleOp, std::move(patterns));
+ }
+};
+} // namespace
+
+void mlir::memref::populateElevateAllocsToGlobalsPatterns(
+ RewritePatternSet &patterns) {
+ patterns.insert<ElevateAllocsToGlobals>(patterns.getContext());
+}
diff --git a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
new file mode 100644
index 0000000000000..3ed66024fd57e
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
@@ -0,0 +1,140 @@
+// RUN: mlir-opt --elevate-allocs-to-globals --split-input-file %s | FileCheck %s
+
+/// Test that a single static memref.alloc is elevated to a memref.global, references
+/// are replaced with memref.get_global, and the associated memref.dealloc is removed.
+
+func.func @single_alloc() -> memref<10xf32> {
+ %0 = memref.alloc() {alignment = 64 : i64} : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
+ return %0 : memref<10xf32>
+}
+
+// CHECK-LABEL: func.func @single_alloc() -> memref<10xf32> {
+// CHECK-NEXT: %[[MEM:.*]] = memref.get_global @global_alloc : memref<10xf32>
+// CHECK-NEXT: return %[[MEM]] : memref<10xf32>
+// CHECK-NOT: memref.dealloc
+// CHECK: memref.global "private" @global_alloc : memref<10xf32> = uninitialized {alignment = 64 : i64}
+
+// -----
+
+/// Test that multiple static memref.alloc ops in the same function are elevated
+/// to global memrefs without symbol name collisions.
+
+func.func @multiple_allocs() -> (memref<10xf32>, memref<20xf32>) {
+ %0 = memref.alloc() : memref<10xf32>
+ %1 = memref.alloc() : memref<20xf32>
+ return %0, %1 : memref<10xf32>, memref<20xf32>
+}
+
+// CHECK-LABEL: func.func @multiple_allocs() -> (memref<10xf32>, memref<20xf32>) {
+// CHECK-DAG: %[[MEM0:.*]] = memref.get_global @global_alloc_0 : memref<10xf32>
+// CHECK-DAG: %[[MEM1:.*]] = memref.get_global @global_alloc : memref<20xf32>
+// CHECK: return %[[MEM0]], %[[MEM1]] : memref<10xf32>, memref<20xf32>
+// CHECK-DAG: memref.global "private" @global_alloc : memref<20xf32> = uninitialized
+// CHECK-DAG: memref.global "private" @global_alloc_0 : memref<10xf32> = uninitialized
+
+// -----
+
+/// Test that a dynamically-shaped memref.alloc is ignored and not elevated to a global.
+
+func.func @dynamic_alloc_ignored(%sz: index) -> memref<?xf32> {
+ %0 = memref.alloc(%sz) : memref<?xf32>
+ return %0 : memref<?xf32>
+}
+
+// CHECK-LABEL: func.func @dynamic_alloc_ignored(
+// CHECK-SAME: %[[SZ:.*]]: index) -> memref<?xf32> {
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[SZ]]) : memref<?xf32>
+// CHECK-NEXT: return %[[MEM]] : memref<?xf32>
+
+// -----
+
+/// Test that a partially dynamic memref.alloc is ignored and not elevated to a global.
+
+func.func @partially_dynamic_alloc_ignored(%sz: index) -> memref<10x?xf32> {
+ %0 = memref.alloc(%sz) : memref<10x?xf32>
+ return %0 : memref<10x?xf32>
+}
+
+// CHECK-LABEL: func.func @partially_dynamic_alloc_ignored(
+// CHECK-SAME: %[[SZ:.*]]: index) -> memref<10x?xf32> {
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[SZ]]) : memref<10x?xf32>
+// CHECK-NEXT: return %[[MEM]] : memref<10x?xf32>
+
+// -----
+
+/// Test that a static memref.alloc inside a loop is ignored and not elevated to a global.
+
+func.func @alloc_in_loop_ignored(%lb: index, %ub: index, %step: index) {
+ scf.for %i = %lb to %ub step %step {
+ %0 = memref.alloc() : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @alloc_in_loop_ignored(
+// CHECK: scf.for
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT: memref.dealloc %[[MEM]] : memref<10xf32>
+
+// -----
+
+/// Test that a static memref.alloc inside control flow (scf.if) is ignored and not
+/// elevated to a global.
+
+func.func @alloc_in_control_flow_ignored(%cond: i1) {
+ scf.if %cond {
+ %0 = memref.alloc() : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @alloc_in_control_flow_ignored(
+// CHECK: scf.if
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT: memref.dealloc %[[MEM]] : memref<10xf32>
+
+// -----
+
+/// Test that a static memref.alloc inside a non-ModuleOp symbol table is
+/// ignored and not elevated to a global.
+
+gpu.module @gpu_mod {
+ gpu.func @kernel() {
+ %0 = memref.alloc() : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
+ gpu.return
+ }
+}
+
+// CHECK-LABEL: gpu.module @gpu_mod
+// CHECK: gpu.func @kernel
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT: memref.dealloc %[[MEM]] : memref<10xf32>
+// CHECK-NOT: memref.global
+
+// -----
+
+/// Test that in a function with allocs both inside and outside of control flow,
+/// only the alloc outside of control flow is elevated to a global.
+
+func.func @mixed_control_flow_allocs(%cond: i1) -> memref<10xf32> {
+ %outside = memref.alloc() : memref<10xf32>
+ scf.if %cond {
+ %inside = memref.alloc() : memref<20xf32>
+ memref.dealloc %inside : memref<20xf32>
+ }
+ return %outside : memref<10xf32>
+}
+
+// CHECK-LABEL: func.func @mixed_control_flow_allocs(
+// CHECK-SAME: %[[COND:.*]]: i1) -> memref<10xf32> {
+// CHECK-NEXT: %[[OUTSIDE:.*]] = memref.get_global @global_alloc : memref<10xf32>
+// CHECK-NEXT: scf.if %[[COND]] {
+// CHECK-NEXT: %[[INSIDE:.*]] = memref.alloc() : memref<20xf32>
+// CHECK-NEXT: memref.dealloc %[[INSIDE]] : memref<20xf32>
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[OUTSIDE]] : memref<10xf32>
+// CHECK: memref.global "private" @global_alloc : memref<10xf32> = uninitialized
``````````
</details>
https://github.com/llvm/llvm-project/pull/211141
More information about the Mlir-commits
mailing list