[Mlir-commits] [mlir] [MLIR][Mem2Reg] Add support for memref.alloca_scope (PR #214221)

[Yue Huang]

https://github.com/AdUhTkJm created https://github.com/llvm/llvm-project/pull/214221

This PR implements `PromotableRegionOpInterface` for `memref.alloca_scope`. This widens the range of programs that can be processed by Mem2Reg pass.

>From 4383e89ec5da6a218195d48900774678506043f0 Mon Sep 17 00:00:00 2001
From: Yue Huang <yh548 at cam.ac.uk>
Date: Wed, 5 Aug 2026 13:07:40 +0100
Subject: [PATCH] [MLIR][Mem2Reg] Add support for memref.alloca_scope

---
 .../mlir/Dialect/MemRef/IR/MemRefOps.td       | 15 +++----
 .../Dialect/MemRef/IR/MemRefMemorySlot.cpp    | 39 ++++++++++++++++++-
 mlir/test/Dialect/MemRef/mem2reg.mlir         | 18 +++++++++
 3 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 9dba4d790d631..4600075dca7e0 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -419,13 +419,14 @@ def MemRef_AllocaOp : AllocLikeOp<"alloca", AutomaticAllocationScopeResource,[
 // AllocaScopeOp
 //===----------------------------------------------------------------------===//
 
-def MemRef_AllocaScopeOp : MemRef_Op<"alloca_scope",
-      [AutomaticAllocationScope,
-       DeclareOpInterfaceMethods<RegionBranchOpInterface, [
-          "getSuccessorInputs"]>,
-       SingleBlockImplicitTerminator<"AllocaScopeReturnOp">,
-       RecursiveMemoryEffects,
-       NoRegionArguments]> {
+def MemRef_AllocaScopeOp
+    : MemRef_Op<"alloca_scope",
+                [AutomaticAllocationScope,
+                 DeclareOpInterfaceMethods<
+                     RegionBranchOpInterface, ["getSuccessorInputs"]>,
+                 DeclareOpInterfaceMethods<PromotableRegionOpInterface>,
+                 SingleBlockImplicitTerminator<"AllocaScopeReturnOp">,
+                 RecursiveMemoryEffects, NoRegionArguments]> {
   let summary = "explicitly delimited scope for stack allocation";
   let description = [{
     The `memref.alloca_scope` operation represents an explicitly-delimited
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index 6748e2cf71804..7782b1ffff239 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -19,8 +19,8 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "mlir/Interfaces/Utils/MemorySlotUtils.h"
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
@@ -325,6 +325,43 @@ struct MemRefDestructurableTypeExternalModel
 
 } // namespace
 
+//===----------------------------------------------------------------------===//
+//  Interfaces for AllocaScopeOp
+//===----------------------------------------------------------------------===//
+bool memref::AllocaScopeOp::isRegionPromotable(const MemorySlot &slot,
+                                               Region *region,
+                                               bool hasValueStores) {
+  return true;
+}
+
+void memref::AllocaScopeOp::setupPromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
+  regionsToProcess.insert({&getRegion(), reachingDef});
+}
+
+Value memref::AllocaScopeOp::finalizePromotion(
+    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
+    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
+    OpBuilder &builder) {
+  if (!hasValueStores)
+    return reachingDef;
+
+  IRRewriter rewriter(builder);
+
+  // Update the return terminator to return the newly defined reaching
+  // definition.
+  memoryslot::updateTerminator(&getRegion().back(), reachingDef,
+                               reachingAtBlockEnd);
+
+  SmallVector<Type> resultTypes(getResultTypes());
+  resultTypes.push_back(slot.elemType);
+
+  Operation *newOp =
+      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
+  return newOp->getResults().back();
+}
+
 //===----------------------------------------------------------------------===//
 //  Register external models
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/MemRef/mem2reg.mlir b/mlir/test/Dialect/MemRef/mem2reg.mlir
index 8f937c4efe75e..932b595f8b30c 100644
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@@ -309,3 +309,21 @@ func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> i32 {
   // CHECK: return %[[RESULT]] : i32
   return %result : i32
 }
+
+// Ensure that AllocaScopeOp does not block mem2reg.
+
+func.func @alloca_scope() -> i32 {
+  %c0 = arith.constant 0 : i32
+  %alloca = memref.alloca() {alignment = 4 : i64} : memref<i32>
+  memref.store %c0, %alloca[] : memref<i32>
+  // CHECK:  %[[RET:.*]] = memref.alloca_scope
+  memref.alloca_scope  {
+    %c1 = arith.constant 1 : i32
+    memref.store %c1, %alloca[] : memref<i32>
+    // CHECK: %[[ONE:.*]] = arith.constant 1
+    // CHECK: memref.alloca_scope.return %[[ONE]]
+  }
+  %value = memref.load %alloca[] : memref<i32>
+  // CHECK: return %[[RET]]
+  return %value : i32
+}