diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -279,6 +279,10 @@ /// Given an operation, return whether this op could to /// allocate an AutomaticAllocationScopeResource static bool isPotentialAutomaticAllocationScope(Operation *op) { + // This op itself doesn't create a stack allocation, + // the inner allocation should be handled separately. + if (op->hasTrait()) + return false; MemoryEffectOpInterface interface = dyn_cast(op); if (!interface) return true; @@ -312,9 +316,11 @@ if (!op->getParentOp()->hasTrait()) { bool hasPotentialAlloca = op->walk([&](Operation *alloc) { + if (alloc == op) + return WalkResult::advance(); if (isPotentialAutomaticAllocationScope(alloc)) return WalkResult::interrupt(); - return WalkResult::skip(); + return WalkResult::advance(); }).wasInterrupted(); if (hasPotentialAlloca) return failure(); diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir --- a/mlir/test/Dialect/MemRef/canonicalize.mlir +++ b/mlir/test/Dialect/MemRef/canonicalize.mlir @@ -643,3 +643,17 @@ // CHECK: }) : () -> () // CHECK: return // CHECK: } + +func @scopeInline(%arg : memref) { + %cnt = "test.count"() : () -> index + "test.region"() ({ + memref.alloca_scope { + memref.store %cnt, %arg[] : memref + } + "test.terminator"() : () -> () + }) : () -> () + return +} + +// CHECK: func @scopeInline +// CHECK-NOT: memref.alloca_scope