[Mlir-commits] [mlir] 1d17915 - [MLIR][MemRef] Ensure alloca_scope is inlined with no allocating ops

[William S. Moses]

Author: William S. Moses
Date: 2022-03-04T11:58:59-05:00
New Revision: 1d1791572cf2172a243ca60bd3ecd7aa2c55e966

URL: https://github.com/llvm/llvm-project/commit/1d1791572cf2172a243ca60bd3ecd7aa2c55e966
DIFF: https://github.com/llvm/llvm-project/commit/1d1791572cf2172a243ca60bd3ecd7aa2c55e966.diff

LOG: [MLIR][MemRef] Ensure alloca_scope is inlined with no allocating ops

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D120841

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Dialect/MemRef/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 7aff36ff2b9db..bf8edc0043488 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -261,7 +261,7 @@ void AllocaScopeOp::getSuccessorRegions(
 
 /// Given an operation, return whether this op is guaranteed to
 /// allocate an AutomaticAllocationScopeResource
-static bool isGuaranteedAutomaticAllocationScope(Operation *op) {
+static bool isGuaranteedAutomaticAllocation(Operation *op) {
   MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op);
   if (!interface)
     return false;
@@ -276,9 +276,15 @@ static bool isGuaranteedAutomaticAllocationScope(Operation *op) {
   return false;
 }
 
-/// Given an operation, return whether this op could to
-/// allocate an AutomaticAllocationScopeResource
-static bool isPotentialAutomaticAllocationScope(Operation *op) {
+/// Given an operation, return whether this op itself could
+/// allocate an AutomaticAllocationScopeResource. Note that
+/// this will not check whether an operation contained within
+/// the op can allocate.
+static bool isOpItselfPotentialAutomaticAllocation(Operation *op) {
+  // This op itself doesn't create a stack allocation,
+  // the inner allocation should be handled separately.
+  if (op->hasTrait<OpTrait::HasRecursiveSideEffects>())
+    return false;
   MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op);
   if (!interface)
     return true;
@@ -312,9 +318,11 @@ struct AllocaScopeInliner : public OpRewritePattern<AllocaScopeOp> {
     if (!op->getParentOp()->hasTrait<OpTrait::AutomaticAllocationScope>()) {
       bool hasPotentialAlloca =
           op->walk([&](Operation *alloc) {
-              if (isPotentialAutomaticAllocationScope(alloc))
+              if (alloc == op)
+                return WalkResult::advance();
+              if (isOpItselfPotentialAutomaticAllocation(alloc))
                 return WalkResult::interrupt();
-              return WalkResult::skip();
+              return WalkResult::advance();
             }).wasInterrupted();
       if (hasPotentialAlloca)
         return failure();
@@ -383,7 +391,7 @@ struct AllocaScopeHoister : public OpRewritePattern<AllocaScopeOp> {
 
     SmallVector<Operation *> toHoist;
     op->walk([&](Operation *alloc) {
-      if (!isGuaranteedAutomaticAllocationScope(alloc))
+      if (!isGuaranteedAutomaticAllocation(alloc))
         return WalkResult::skip();
 
       // If any operand is not defined before the location of

diff  --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index 12c5a529e901d..be3ac32cbfbda 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -643,3 +643,17 @@ func @scopeMerge4() {
 // CHECK:     }) : () -> ()
 // CHECK:     return
 // CHECK:   }
+
+func @scopeInline(%arg : memref<index>) {
+  %cnt = "test.count"() : () -> index
+  "test.region"() ({
+    memref.alloca_scope {
+      memref.store %cnt, %arg[] : memref<index>
+    }
+    "test.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK:   func @scopeInline
+// CHECK-NOT:  memref.alloca_scope