[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)
Bhavesh M
llvmlistbot at llvm.org
Thu Jul 23 13:54:16 PDT 2026
https://github.com/beamandala updated https://github.com/llvm/llvm-project/pull/211141
>From 4ff9cf5e7abbc47b3f1ab27b44f821eef44763e4 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Tue, 21 Jul 2026 16:07:42 -0700
Subject: [PATCH 1/7] Elevate allocs to globals pass
---
.../mlir/Dialect/MemRef/Transforms/Passes.td | 10 ++
.../Dialect/MemRef/Transforms/Transforms.h | 3 +
.../Dialect/MemRef/Transforms/CMakeLists.txt | 1 +
.../Transforms/ElevateAllocsToGlobals.cpp | 112 ++++++++++++++++++
.../MemRef/elevate-allocs-to-globals.mlir | 56 +++++++++
5 files changed, 182 insertions(+)
create mode 100644 mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
create mode 100644 mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index c6be600247696..4e6ead48cbff2 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -340,4 +340,14 @@ def FlattenMemrefsPass : Pass<"flatten-memref"> {
];
}
+def ElevateAllocsToGlobalsPass : Pass<"elevate-allocs-to-globals", "ModuleOp"> {
+ let summary = "Elevate allocs to globals";
+ let description = [{
+
+ }];
+ 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..017b23b810dd3
--- /dev/null
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -0,0 +1,112 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/Interfaces/LoopLikeInterface.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.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 {
+
+// Checks if 'op' is contained inside any branching or looping structure
+static bool isInsideControlFlow(mlir::Operation *op) {
+ if (mlir::getEnclosingRepetitiveRegion(op) != nullptr)
+ return true;
+
+ if (op->getParentOfType<mlir::LoopLikeOpInterface>())
+ return true;
+
+ if (auto regionParent = op->getParentOfType<mlir::RegionBranchOpInterface>())
+ return true;
+
+ return false;
+}
+
+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
+ if (!memrefType.hasStaticShape() || !allocOp.getDynamicSizes().empty())
+ return failure();
+
+ auto loopParent = allocOp->getParentOfType<mlir::LoopLikeOpInterface>();
+ if (loopParent != nullptr || isInsideControlFlow(allocOp))
+ return failure();
+
+ memref::GlobalOp globalOp;
+ {
+ Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(allocOp);
+
+ SymbolTable symbolTable(symbolTableOp);
+
+ OpBuilder builder(rewriter.getContext());
+ StringAttr globalName = rewriter.getStringAttr("global_alloc");
+ globalOp = memref::GlobalOp::create(builder, allocOp.getLoc(), globalName,
+ rewriter.getStringAttr("private"),
+ memrefType, rewriter.getUnitAttr(),
+ false, allocOp.getAlignmentAttr());
+
+ symbolTable.insert(globalOp);
+ }
+
+ 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);
+
+ 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());
+}
\ No newline at end of file
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..dd65cd12989cc
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
@@ -0,0 +1,56 @@
+// 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,
+/// replaced with memref.get_global, alignment attribute is preserved, and
+/// associated memref.dealloc is removed.
+
+func.func @single_alloc(%val: f32, %idx: index) {
+ %0 = memref.alloc() {alignment = 64 : i64} : memref<10x20xf32>
+ memref.store %val, %0[%idx, %idx] : memref<10x20xf32>
+ memref.dealloc %0 : memref<10x20xf32>
+ return
+}
+
+// CHECK-LABEL: func.func @single_alloc(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
+// CHECK-NEXT: %[[MEM:.*]] = memref.get_global @global_alloc : memref<10x20xf32>
+// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]], %[[ARG1]]] : memref<10x20xf32>
+// CHECK-NEXT: return
+// CHECK-NOT: memref.dealloc
+// CHECK: memref.global "private" @global_alloc : memref<10x20xf32> = 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(%val: f32, %val_i32: i32, %idx: index) {
+ %0 = memref.alloc() : memref<10xf32>
+ %1 = memref.alloc() : memref<20xi32>
+ memref.store %val, %0[%idx] : memref<10xf32>
+ memref.store %val_i32, %1[%idx] : memref<20xi32>
+ return
+}
+
+// CHECK-LABEL: func.func @multiple_allocs(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: index) {
+// CHECK-DAG: %[[MEM0:.*]] = memref.get_global @global_alloc_0 : memref<10xf32>
+// CHECK-DAG: %[[MEM1:.*]] = memref.get_global @global_alloc : memref<20xi32>
+// CHECK: memref.store %[[ARG0]], %[[MEM0]][%[[ARG2]]] : memref<10xf32>
+// CHECK: memref.store %[[ARG1]], %[[MEM1]][%[[ARG2]]] : memref<20xi32>
+// CHECK-DAG: memref.global "private" @global_alloc : memref<20xi32> = 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(%val: f32, %sz: index) {
+ %0 = memref.alloc(%sz) : memref<?xf32>
+ memref.store %val, %0[%sz] : memref<?xf32>
+ return
+}
+
+// CHECK-LABEL: func.func @dynamic_alloc_ignored(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
+// CHECK: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<?xf32>
+// CHECK: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<?xf32>
>From b0f3250cd4e5a865afc304d8a9a04674fe33cd25 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Tue, 21 Jul 2026 17:48:50 -0700
Subject: [PATCH 2/7] Simplify and add tests
---
.../Transforms/ElevateAllocsToGlobals.cpp | 2 +-
.../MemRef/elevate-allocs-to-globals.mlir | 88 +++++++++++++++----
2 files changed, 70 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
index 017b23b810dd3..2dbdeb00c5354 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -109,4 +109,4 @@ struct ElevateAllocsToGlobalsPass
void mlir::memref::populateElevateAllocsToGlobalsPatterns(
RewritePatternSet &patterns) {
patterns.insert<ElevateAllocsToGlobals>(patterns.getContext());
-}
\ No newline at end of file
+}
diff --git a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
index dd65cd12989cc..43477960cf7be 100644
--- a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
+++ b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
@@ -1,49 +1,49 @@
// 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,
-/// replaced with memref.get_global, alignment attribute is preserved, and
-/// associated memref.dealloc is removed.
+/// 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(%val: f32, %idx: index) {
- %0 = memref.alloc() {alignment = 64 : i64} : memref<10x20xf32>
- memref.store %val, %0[%idx, %idx] : memref<10x20xf32>
- memref.dealloc %0 : memref<10x20xf32>
+ %0 = memref.alloc() {alignment = 64 : i64} : memref<10xf32>
+ memref.store %val, %0[%idx] : memref<10xf32>
+ memref.dealloc %0 : memref<10xf32>
return
}
// CHECK-LABEL: func.func @single_alloc(
// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK-NEXT: %[[MEM:.*]] = memref.get_global @global_alloc : memref<10x20xf32>
-// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]], %[[ARG1]]] : memref<10x20xf32>
+// CHECK-NEXT: %[[MEM:.*]] = memref.get_global @global_alloc : memref<10xf32>
+// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<10xf32>
// CHECK-NEXT: return
// CHECK-NOT: memref.dealloc
-// CHECK: memref.global "private" @global_alloc : memref<10x20xf32> = uninitialized {alignment = 64 : i64}
+// 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(%val: f32, %val_i32: i32, %idx: index) {
+func.func @multiple_allocs(%val: f32, %idx: index) {
%0 = memref.alloc() : memref<10xf32>
- %1 = memref.alloc() : memref<20xi32>
+ %1 = memref.alloc() : memref<20xf32>
memref.store %val, %0[%idx] : memref<10xf32>
- memref.store %val_i32, %1[%idx] : memref<20xi32>
+ memref.store %val, %1[%idx] : memref<20xf32>
return
}
// CHECK-LABEL: func.func @multiple_allocs(
-// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: index) {
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
// CHECK-DAG: %[[MEM0:.*]] = memref.get_global @global_alloc_0 : memref<10xf32>
-// CHECK-DAG: %[[MEM1:.*]] = memref.get_global @global_alloc : memref<20xi32>
-// CHECK: memref.store %[[ARG0]], %[[MEM0]][%[[ARG2]]] : memref<10xf32>
-// CHECK: memref.store %[[ARG1]], %[[MEM1]][%[[ARG2]]] : memref<20xi32>
-// CHECK-DAG: memref.global "private" @global_alloc : memref<20xi32> = uninitialized
+// CHECK-DAG: %[[MEM1:.*]] = memref.get_global @global_alloc : memref<20xf32>
+// CHECK: memref.store %[[ARG0]], %[[MEM0]][%[[ARG1]]] : memref<10xf32>
+// CHECK: memref.store %[[ARG0]], %[[MEM1]][%[[ARG1]]] : 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(%val: f32, %sz: index) {
%0 = memref.alloc(%sz) : memref<?xf32>
memref.store %val, %0[%sz] : memref<?xf32>
@@ -52,5 +52,55 @@ func.func @dynamic_alloc_ignored(%val: f32, %sz: index) {
// CHECK-LABEL: func.func @dynamic_alloc_ignored(
// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<?xf32>
-// CHECK: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<?xf32>
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<?xf32>
+// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<?xf32>
+
+// -----
+
+/// Test that a partially dynamic memref.alloc is ignored and not elevated to a global.
+
+func.func @partially_dynamic_alloc_ignored(%val: f32, %sz: index) {
+ %0 = memref.alloc(%sz) : memref<10x?xf32>
+ memref.store %val, %0[%sz, %sz] : memref<10x?xf32>
+ return
+}
+
+// CHECK-LABEL: func.func @partially_dynamic_alloc_ignored(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<10x?xf32>
+// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]], %[[ARG1]]] : 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, %val: f32, %idx: index) {
+ scf.for %i = %lb to %ub step %step {
+ %0 = memref.alloc() : memref<10xf32>
+ memref.store %val, %0[%idx] : memref<10xf32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @alloc_in_loop_ignored(
+// CHECK: scf.for
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT: memref.store %{{.*}}, %[[MEM]]
+
+// -----
+
+/// 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, %val: f32, %idx: index) {
+ scf.if %cond {
+ %0 = memref.alloc() : memref<10xf32>
+ memref.store %val, %0[%idx] : 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.store %{{.*}}, %[[MEM]]
>From d7f53ae83635215983dd626720eab2f7db5e3593 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 22 Jul 2026 15:29:11 -0700
Subject: [PATCH 3/7] Pass description
---
.../mlir/Dialect/MemRef/Transforms/Passes.td | 24 +++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index 4e6ead48cbff2..a0a4be7318889 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -343,7 +343,31 @@ 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"
>From 8b4d7a5e2cccc7b928ce05f986ff40c2e5afb6c7 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 22 Jul 2026 15:36:18 -0700
Subject: [PATCH 4/7] Clean up helper func
---
.../Transforms/ElevateAllocsToGlobals.cpp | 20 ++++++-------------
1 file changed, 6 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
index 2dbdeb00c5354..5a079ac8c2b63 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -28,18 +28,11 @@ using namespace mlir;
namespace {
-// Checks if 'op' is contained inside any branching or looping structure
-static bool isInsideControlFlow(mlir::Operation *op) {
- if (mlir::getEnclosingRepetitiveRegion(op) != nullptr)
- return true;
-
- if (op->getParentOfType<mlir::LoopLikeOpInterface>())
- return true;
-
- if (auto regionParent = op->getParentOfType<mlir::RegionBranchOpInterface>())
- return true;
-
- return false;
+// Checks if 'op' is contained inside any branching or looping structure.
+static bool isInsideControlFlow(Operation *op) {
+ return getEnclosingRepetitiveRegion(op) ||
+ op->getParentOfType<LoopLikeOpInterface>() ||
+ op->getParentOfType<RegionBranchOpInterface>();
}
struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
@@ -54,8 +47,7 @@ struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
if (!memrefType.hasStaticShape() || !allocOp.getDynamicSizes().empty())
return failure();
- auto loopParent = allocOp->getParentOfType<mlir::LoopLikeOpInterface>();
- if (loopParent != nullptr || isInsideControlFlow(allocOp))
+ if (isInsideControlFlow(allocOp))
return failure();
memref::GlobalOp globalOp;
>From eb611a58fb3a322d0236d71410bc654ffd1af43d Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 22 Jul 2026 15:48:22 -0700
Subject: [PATCH 5/7] Comments
---
.../Transforms/ElevateAllocsToGlobals.cpp | 29 ++++++++++++++++---
1 file changed, 25 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
index 5a079ac8c2b63..ec1130e96cb75 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -28,32 +28,51 @@ using namespace mlir;
namespace {
-// Checks if 'op' is contained inside any branching or looping structure.
+/// 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
+ // `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 (e.g.
+ // module).
memref::GlobalOp globalOp;
{
Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(allocOp);
-
SymbolTable symbolTable(symbolTableOp);
OpBuilder builder(rewriter.getContext());
@@ -66,6 +85,7 @@ struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
symbolTable.insert(globalOp);
}
+ // Remove any `memref.dealloc` operations using this allocation
SmallVector<Operation *> deallocsToDelete;
for (OpOperand &use : allocOp.getResult().getUses()) {
Operation *user = use.getOwner();
@@ -75,6 +95,7 @@ struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
for (Operation *dealloc : deallocsToDelete)
rewriter.eraseOp(dealloc);
+ // Replace the original `memref.alloc` with `memref.get_global`.
rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(allocOp, memrefType,
globalOp.getName());
>From ca1e6c5c6e3c03923b7af10024ce4b8d10bf2ea9 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Thu, 23 Jul 2026 13:06:42 -0700
Subject: [PATCH 6/7] restrict MemRef ElevateAllocsToGlobals to ModuleOps and
add tests
---
.../Transforms/ElevateAllocsToGlobals.cpp | 33 +++++++-------
.../MemRef/elevate-allocs-to-globals.mlir | 45 +++++++++++++++++++
2 files changed, 62 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
index ec1130e96cb75..907038b545c62 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
@@ -13,8 +13,10 @@
#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 {
@@ -68,22 +70,21 @@ struct ElevateAllocsToGlobals : public OpRewritePattern<memref::AllocOp> {
if (isInsideControlFlow(allocOp))
return failure();
- // Create the global variable at the nearest enclosing symbol table (e.g.
- // module).
- memref::GlobalOp globalOp;
- {
- Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(allocOp);
- SymbolTable symbolTable(symbolTableOp);
-
- OpBuilder builder(rewriter.getContext());
- StringAttr globalName = rewriter.getStringAttr("global_alloc");
- globalOp = memref::GlobalOp::create(builder, allocOp.getLoc(), globalName,
- rewriter.getStringAttr("private"),
- memrefType, rewriter.getUnitAttr(),
- false, allocOp.getAlignmentAttr());
-
- symbolTable.insert(globalOp);
- }
+ // 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;
diff --git a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
index 43477960cf7be..b91e5f92c31e1 100644
--- a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
+++ b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
@@ -104,3 +104,48 @@ func.func @alloc_in_control_flow_ignored(%cond: i1, %val: f32, %idx: index) {
// CHECK: scf.if
// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
// CHECK-NEXT: memref.store %{{.*}}, %[[MEM]]
+
+// -----
+
+/// 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, %val: f32, %idx: index) {
+ %outside = memref.alloc() : memref<10xf32>
+ memref.store %val, %outside[%idx] : memref<10xf32>
+ scf.if %cond {
+ %inside = memref.alloc() : memref<20xf32>
+ memref.store %val, %inside[%idx] : memref<20xf32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @mixed_control_flow_allocs(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[VAL:.*]]: f32, %[[IDX:.*]]: index) {
+// CHECK-NEXT: %[[OUTSIDE:.*]] = memref.get_global @global_alloc : memref<10xf32>
+// CHECK-NEXT: memref.store %[[VAL]], %[[OUTSIDE]][%[[IDX]]] : memref<10xf32>
+// CHECK-NEXT: scf.if %[[COND]] {
+// CHECK-NEXT: %[[INSIDE:.*]] = memref.alloc() : memref<20xf32>
+// CHECK-NEXT: memref.store %[[VAL]], %[[INSIDE]][%[[IDX]]] : memref<20xf32>
+// CHECK-NEXT: }
+// CHECK-NEXT: return
+// CHECK: memref.global "private" @global_alloc : memref<10xf32> = uninitialized
>From 694b3704b165038b349ff3f4700abe27fb9c833b Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Thu, 23 Jul 2026 13:16:22 -0700
Subject: [PATCH 7/7] Update test cases
---
.../MemRef/elevate-allocs-to-globals.mlir | 111 ++++++++----------
1 file changed, 50 insertions(+), 61 deletions(-)
diff --git a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
index b91e5f92c31e1..3ed66024fd57e 100644
--- a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
+++ b/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
@@ -3,107 +3,98 @@
/// 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(%val: f32, %idx: index) {
+func.func @single_alloc() -> memref<10xf32> {
%0 = memref.alloc() {alignment = 64 : i64} : memref<10xf32>
- memref.store %val, %0[%idx] : memref<10xf32>
memref.dealloc %0 : memref<10xf32>
- return
+ return %0 : memref<10xf32>
}
-// CHECK-LABEL: func.func @single_alloc(
-// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK-NEXT: %[[MEM:.*]] = memref.get_global @global_alloc : memref<10xf32>
-// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<10xf32>
-// CHECK-NEXT: return
-// CHECK-NOT: memref.dealloc
-// CHECK: memref.global "private" @global_alloc : memref<10xf32> = uninitialized {alignment = 64 : i64}
+// 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(%val: f32, %idx: index) {
+func.func @multiple_allocs() -> (memref<10xf32>, memref<20xf32>) {
%0 = memref.alloc() : memref<10xf32>
%1 = memref.alloc() : memref<20xf32>
- memref.store %val, %0[%idx] : memref<10xf32>
- memref.store %val, %1[%idx] : memref<20xf32>
- return
+ return %0, %1 : memref<10xf32>, memref<20xf32>
}
-// CHECK-LABEL: func.func @multiple_allocs(
-// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK-DAG: %[[MEM0:.*]] = memref.get_global @global_alloc_0 : memref<10xf32>
-// CHECK-DAG: %[[MEM1:.*]] = memref.get_global @global_alloc : memref<20xf32>
-// CHECK: memref.store %[[ARG0]], %[[MEM0]][%[[ARG1]]] : memref<10xf32>
-// CHECK: memref.store %[[ARG0]], %[[MEM1]][%[[ARG1]]] : memref<20xf32>
-// CHECK-DAG: memref.global "private" @global_alloc : memref<20xf32> = uninitialized
-// CHECK-DAG: memref.global "private" @global_alloc_0 : memref<10xf32> = uninitialized
+// 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(%val: f32, %sz: index) {
+func.func @dynamic_alloc_ignored(%sz: index) -> memref<?xf32> {
%0 = memref.alloc(%sz) : memref<?xf32>
- memref.store %val, %0[%sz] : memref<?xf32>
- return
+ return %0 : memref<?xf32>
}
// CHECK-LABEL: func.func @dynamic_alloc_ignored(
-// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<?xf32>
-// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]]] : memref<?xf32>
+// 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(%val: f32, %sz: index) {
+func.func @partially_dynamic_alloc_ignored(%sz: index) -> memref<10x?xf32> {
%0 = memref.alloc(%sz) : memref<10x?xf32>
- memref.store %val, %0[%sz, %sz] : memref<10x?xf32>
- return
+ return %0 : memref<10x?xf32>
}
// CHECK-LABEL: func.func @partially_dynamic_alloc_ignored(
-// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: index) {
-// CHECK-NEXT: %[[MEM:.*]] = memref.alloc(%[[ARG1]]) : memref<10x?xf32>
-// CHECK-NEXT: memref.store %[[ARG0]], %[[MEM]][%[[ARG1]], %[[ARG1]]] : memref<10x?xf32>
+// 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, %val: f32, %idx: index) {
+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.store %val, %0[%idx] : 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.store %{{.*}}, %[[MEM]]
+// 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, %val: f32, %idx: index) {
+func.func @alloc_in_control_flow_ignored(%cond: i1) {
scf.if %cond {
%0 = memref.alloc() : memref<10xf32>
- memref.store %val, %0[%idx] : 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.store %{{.*}}, %[[MEM]]
+// CHECK: scf.if
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<10xf32>
+// CHECK-NEXT: memref.dealloc %[[MEM]] : memref<10xf32>
// -----
@@ -119,33 +110,31 @@ gpu.module @gpu_mod {
}
// 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
+// 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, %val: f32, %idx: index) {
+func.func @mixed_control_flow_allocs(%cond: i1) -> memref<10xf32> {
%outside = memref.alloc() : memref<10xf32>
- memref.store %val, %outside[%idx] : memref<10xf32>
scf.if %cond {
%inside = memref.alloc() : memref<20xf32>
- memref.store %val, %inside[%idx] : memref<20xf32>
+ memref.dealloc %inside : memref<20xf32>
}
- return
+ return %outside : memref<10xf32>
}
// CHECK-LABEL: func.func @mixed_control_flow_allocs(
-// CHECK-SAME: %[[COND:.*]]: i1, %[[VAL:.*]]: f32, %[[IDX:.*]]: index) {
-// CHECK-NEXT: %[[OUTSIDE:.*]] = memref.get_global @global_alloc : memref<10xf32>
-// CHECK-NEXT: memref.store %[[VAL]], %[[OUTSIDE]][%[[IDX]]] : memref<10xf32>
-// CHECK-NEXT: scf.if %[[COND]] {
-// CHECK-NEXT: %[[INSIDE:.*]] = memref.alloc() : memref<20xf32>
-// CHECK-NEXT: memref.store %[[VAL]], %[[INSIDE]][%[[IDX]]] : memref<20xf32>
-// CHECK-NEXT: }
-// CHECK-NEXT: return
-// CHECK: memref.global "private" @global_alloc : memref<10xf32> = uninitialized
+// 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
More information about the Mlir-commits
mailing list