[Mlir-commits] [mlir] [mlir] [memref] Elevate `AllocOp`s to `GlobalOp`s pass (PR #211141)
Bhavesh M
llvmlistbot at llvm.org
Mon Aug 3 13:50:29 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/9] 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/9] 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/9] 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/9] 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/9] 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/9] 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/9] 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
>From 84d1239622c0a48f0e45c3babc44d308c43bd6c9 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 29 Jul 2026 12:27:56 -0700
Subject: [PATCH 8/9] Remove elevateallocs pass
---
.../mlir/Dialect/MemRef/Transforms/Passes.td | 34 -----
.../Dialect/MemRef/Transforms/Transforms.h | 3 -
.../Dialect/MemRef/Transforms/CMakeLists.txt | 1 -
.../Transforms/ElevateAllocsToGlobals.cpp | 126 ----------------
.../MemRef/elevate-allocs-to-globals.mlir | 140 ------------------
5 files changed, 304 deletions(-)
delete mode 100644 mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
delete 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 a0a4be7318889..c6be600247696 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -340,38 +340,4 @@ 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 6960b8ef00f06..720677455ae5d 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
@@ -14,7 +14,6 @@
#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"
@@ -166,8 +165,6 @@ 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 7d43e66989bc4..1a8b03dabcfb7 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
@@ -17,7 +17,6 @@ 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
deleted file mode 100644
index 907038b545c62..0000000000000
--- a/mlir/lib/Dialect/MemRef/Transforms/ElevateAllocsToGlobals.cpp
+++ /dev/null
@@ -1,126 +0,0 @@
-//===----------------------------------------------------------------------===//
-//
-// 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
deleted file mode 100644
index 3ed66024fd57e..0000000000000
--- a/mlir/test/Dialect/MemRef/elevate-allocs-to-globals.mlir
+++ /dev/null
@@ -1,140 +0,0 @@
-// 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
>From 7d7f9cc7cf4559ddb0a66e65c788a189f82a6cb8 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Mon, 3 Aug 2026 13:49:48 -0700
Subject: [PATCH 9/9] MemRefAllocToGlobalOp transform
---
.../MemRef/TransformOps/MemRefTransformOps.td | 73 +++++++++-
.../TransformOps/MemRefTransformOps.cpp | 129 ++++++++++++++----
mlir/test/Dialect/MemRef/transform-ops.mlir | 66 +++++++++
3 files changed, 242 insertions(+), 26 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td b/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td
index f4694a30a8a12..12f465b11c286 100644
--- a/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td
@@ -193,7 +193,9 @@ def MemRefAllocaToGlobalOp :
#### Return modes
- Succeeds always. The returned handles refer to the `memref.get_global` and
+ Succeeds if all payload operations are statically shaped allocations.
+ Fails with a silenceable error if any allocation has a dynamic shape.
+ The returned handles refer to the `memref.get_global` and
`memref.global` ops that were inserted by the transformation.
}];
@@ -206,6 +208,75 @@ def MemRefAllocaToGlobalOp :
}];
}
+def MemRefAllocToGlobalOp :
+ Op<Transform_Dialect, "memref.alloc_to_global",
+ [TransformOpInterface,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+ DeclareOpInterfaceMethods<TransformOpInterface>]> {
+ let description = [{
+ Inserts a new `memref.global` for each provided `memref.alloc` into the
+ nearest symbol table (e.g., a `builtin.module`) and replaces it with a
+ `memref.get_global`. Any `memref.dealloc` operations referencing the
+ allocation are also removed, since the memory is no longer dynamically
+ allocated.
+
+ This transformation assumes that the allocation does not escape the current
+ container (e.g., it is not returned from the function or passed to another
+ function that deallocates it).
+
+ #### Example
+
+ Consider the following transform op:
+
+ ```mlir
+ %get_global, %global =
+ transform.memref.alloc_to_global %alloc
+ : (!transform.op<"memref.alloc">)
+ -> (!transform.any_op, !transform.any_op)
+ ```
+
+ and the following input payload:
+
+ ```mlir
+ module {
+ func.func @func() {
+ %alloc = memref.alloc() : memref<2x32xf32>
+ // usages of %alloc...
+ memref.dealloc %alloc : memref<2x32xf32>
+ }
+ }
+ ```
+
+ then applying the transform op to the payload would result in the following
+ output IR:
+
+ ```mlir
+ module {
+ memref.global "private" @alloc : memref<2x32xf32>
+ func.func @func() {
+ %alloc = memref.get_global @alloc : memref<2x32xf32>
+ // usages of %alloc...
+ }
+ }
+ ```
+
+ #### Return modes
+
+ Succeeds if all payload operations are statically shaped allocations.
+ Fails with a silenceable error if any allocation has a dynamic shape.
+ The returned handles refer to the `memref.get_global` and
+ `memref.global` ops that were inserted by the transformation.
+ }];
+
+ let arguments = (ins Transform_MemRefAllocOp:$alloc);
+ let results = (outs TransformHandleTypeInterface:$getGlobal,
+ TransformHandleTypeInterface:$global);
+
+ let assemblyFormat = [{
+ $alloc attr-dict `:` functional-type(operands, results)
+ }];
+}
+
def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
DeclareOpInterfaceMethods<TransformOpInterface>]> {
diff --git a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
index 95eb2a9a95bc1..fef4656696481 100644
--- a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
+++ b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
@@ -125,6 +125,64 @@ void transform::ApplyResolveRankedShapedTypeResultDimsPatternsOp::
memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
}
+//===----------------------------------------------------------------------===//
+// Alloc and alloca to global utilities
+//===----------------------------------------------------------------------===//
+
+/// Converts an allocation operation (`memref.alloca` or `memref.alloc`) to a
+/// `memref.global` operation in the nearest symbol table, and replaces the
+/// allocation with a `memref.get_global` operation. Any `memref.dealloc`
+/// operations referencing the allocation are erased.
+template <typename AllocLikeOp>
+static DiagnosedSilenceableFailure
+allocLikeToGlobal(transform::TransformRewriter &rewriter,
+ AllocLikeOp allocLikeOp, StringRef globalName,
+ memref::GlobalOp &globalOp,
+ memref::GetGlobalOp &getGlobalOp) {
+ MemRefType memrefType = allocLikeOp.getType();
+ if (!memrefType.hasStaticShape()) {
+ return emitSilenceableFailure(allocLikeOp->getLoc())
+ << "global ops require statically shaped memrefs, but got "
+ << memrefType;
+ }
+
+ MLIRContext *ctx = rewriter.getContext();
+ Location loc = allocLikeOp->getLoc();
+
+ // Find nearest symbol table.
+ Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(allocLikeOp);
+ assert(symbolTableOp && "expected payload to be in symbol table");
+ SymbolTable symbolTable(symbolTableOp);
+
+ // Insert a `memref.global` into the symbol table.
+ Type resultType = allocLikeOp.getResult().getType();
+ OpBuilder builder(rewriter.getContext());
+ // TODO: Add a better builder for this.
+ globalOp = memref::GlobalOp::create(
+ builder, loc, StringAttr::get(ctx, globalName),
+ StringAttr::get(ctx, "private"), TypeAttr::get(resultType), Attribute{},
+ UnitAttr{}, IntegerAttr{});
+ symbolTable.insert(globalOp);
+
+ // Remove any `memref.dealloc` operations referencing this allocation.
+ // We assume that the allocation does not escape the current container
+ // (e.g., via return or interprocedural function calls), so any deallocation
+ // is a direct user of the allocation. This scopes complexity and avoids
+ // the need for interprocedural escape analysis.
+ for (Operation *user : llvm::make_early_inc_range(allocLikeOp->getUsers())) {
+ if (auto dealloc = dyn_cast<memref::DeallocOp>(user))
+ rewriter.eraseOp(dealloc);
+ }
+
+ // Replace the allocation with a `memref.get_global` accessing the
+ // global symbol inserted above.
+ rewriter.setInsertionPoint(allocLikeOp);
+ getGlobalOp = rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(
+ allocLikeOp, globalOp.getType(), globalOp.getName());
+
+ return DiagnosedSilenceableFailure::success();
+}
+
//===----------------------------------------------------------------------===//
// AllocaToGlobalOp
//===----------------------------------------------------------------------===//
@@ -141,32 +199,12 @@ transform::MemRefAllocaToGlobalOp::apply(transform::TransformRewriter &rewriter,
// Transform `memref.alloca`s.
for (auto *op : allocaOps) {
auto alloca = cast<memref::AllocaOp>(op);
- MLIRContext *ctx = rewriter.getContext();
- Location loc = alloca->getLoc();
-
memref::GlobalOp globalOp;
- {
- // Find nearest symbol table.
- Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(op);
- assert(symbolTableOp && "expected alloca payload to be in symbol table");
- SymbolTable symbolTable(symbolTableOp);
-
- // Insert a `memref.global` into the symbol table.
- Type resultType = alloca.getResult().getType();
- OpBuilder builder(rewriter.getContext());
- // TODO: Add a better builder for this.
- globalOp = memref::GlobalOp::create(
- builder, loc, StringAttr::get(ctx, "alloca"),
- StringAttr::get(ctx, "private"), TypeAttr::get(resultType),
- Attribute{}, UnitAttr{}, IntegerAttr{});
- symbolTable.insert(globalOp);
- }
-
- // Replace the `memref.alloca` with a `memref.get_global` accessing the
- // global symbol inserted above.
- rewriter.setInsertionPoint(alloca);
- auto getGlobalOp = rewriter.replaceOpWithNewOp<memref::GetGlobalOp>(
- alloca, globalOp.getType(), globalOp.getName());
+ memref::GetGlobalOp getGlobalOp;
+ DiagnosedSilenceableFailure diag =
+ allocLikeToGlobal(rewriter, alloca, "alloca", globalOp, getGlobalOp);
+ if (!diag.succeeded())
+ return diag;
globalOps.push_back(globalOp);
getGlobalOps.push_back(getGlobalOp);
@@ -186,6 +224,47 @@ void transform::MemRefAllocaToGlobalOp::getEffects(
modifiesPayload(effects);
}
+//===----------------------------------------------------------------------===//
+// AllocToGlobalOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::MemRefAllocToGlobalOp::apply(transform::TransformRewriter &rewriter,
+ transform::TransformResults &results,
+ transform::TransformState &state) {
+ auto allocOps = state.getPayloadOps(getAlloc());
+
+ SmallVector<memref::GlobalOp> globalOps;
+ SmallVector<memref::GetGlobalOp> getGlobalOps;
+
+ // Transform `memref.alloc`s.
+ for (auto *op : allocOps) {
+ auto alloc = cast<memref::AllocOp>(op);
+ memref::GlobalOp globalOp;
+ memref::GetGlobalOp getGlobalOp;
+ DiagnosedSilenceableFailure diag =
+ allocLikeToGlobal(rewriter, alloc, "alloc", globalOp, getGlobalOp);
+ if (!diag.succeeded())
+ return diag;
+
+ globalOps.push_back(globalOp);
+ getGlobalOps.push_back(getGlobalOp);
+ }
+
+ // Assemble results.
+ results.set(cast<OpResult>(getGlobal()), globalOps);
+ results.set(cast<OpResult>(getGetGlobal()), getGlobalOps);
+
+ return DiagnosedSilenceableFailure::success();
+}
+
+void transform::MemRefAllocToGlobalOp::getEffects(
+ SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+ producesHandle(getOperation()->getOpResults(), effects);
+ consumesHandle(getAllocMutable(), effects);
+ modifiesPayload(effects);
+}
+
//===----------------------------------------------------------------------===//
// MemRefMultiBufferOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/MemRef/transform-ops.mlir b/mlir/test/Dialect/MemRef/transform-ops.mlir
index 7fc84d419f18d..8e14c0a2c7e01 100644
--- a/mlir/test/Dialect/MemRef/transform-ops.mlir
+++ b/mlir/test/Dialect/MemRef/transform-ops.mlir
@@ -33,6 +33,72 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-DAG: memref.global "private" @[[ALLOC0:alloc.*]] : memref<2xf32>
+// CHECK-DAG: memref.global "private" @[[ALLOC1:alloc.*]] : memref<2xf32>
+
+// CHECK-DAG: func.func @func_alloc()
+func.func @func_alloc() {
+ // CHECK-DAG: %[[MR0:.*]] = memref.get_global @[[ALLOC0]] : memref<2xf32>
+ // CHECK-DAG: %[[MR1:.*]] = memref.get_global @[[ALLOC1]] : memref<2xf32>
+ // CHECK-NOT: memref.dealloc
+ %mr0 = memref.alloc() : memref<2xf32>
+ %mr1 = memref.alloc() : memref<2xf32>
+ memref.dealloc %mr0 : memref<2xf32>
+ memref.dealloc %mr1 : memref<2xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %alloc = transform.structured.match ops{["memref.alloc"]} in %arg0
+ : (!transform.any_op) -> !transform.op<"memref.alloc">
+ %get_global, %global = transform.memref.alloc_to_global %alloc
+ : (!transform.op<"memref.alloc">)
+ -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @alloc_to_global_dynamic_shape(%arg0: index) {
+ // expected-error @below {{global ops require statically shaped memrefs, but got 'memref<?xf32>'}}
+ %alloc = memref.alloc(%arg0) : memref<?xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %alloc = transform.structured.match ops{["memref.alloc"]} in %arg0
+ : (!transform.any_op) -> !transform.op<"memref.alloc">
+ %get_global, %global = transform.memref.alloc_to_global %alloc
+ : (!transform.op<"memref.alloc">)
+ -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @alloca_to_global_dynamic_shape(%arg0: index) {
+ // expected-error @below {{global ops require statically shaped memrefs, but got 'memref<?xf32>'}}
+ %alloca = memref.alloca(%arg0) : memref<?xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %alloca = transform.structured.match ops{["memref.alloca"]} in %arg0
+ : (!transform.any_op) -> !transform.op<"memref.alloca">
+ %get_global, %global = transform.memref.alloca_to_global %alloca
+ : (!transform.op<"memref.alloca">)
+ -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> ((d0 floordiv 4) mod 2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
More information about the Mlir-commits
mailing list