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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Yue Huang (AdUhTkJm)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/214221.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+8-7) 
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp (+38-1) 
- (modified) mlir/test/Dialect/MemRef/mem2reg.mlir (+18) 


``````````diff
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
+}

``````````

</details>


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