[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)
Bhavesh M
llvmlistbot at llvm.org
Tue Jul 21 17:49:04 PDT 2026
https://github.com/beamandala updated https://github.com/llvm/llvm-project/pull/211141
>From 829c25c037707f17d8c98fcb15c041914a26a46d 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/2] 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 39d4e08842d4dafd659020418859af23193db91c 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/2] 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]]
More information about the Mlir-commits
mailing list