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

[wsmoses]

Add recursive side effects

[ftynse]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

278–287

I find the name confusing and the documentation wrong given the current body of the function. It does not check whether an operation is an allocation _scope_, but whether an operation may be allocating. With this change, it only checks whether the operation _itself_ allocates.

[wsmoses]

Change name/docs

[ftynse]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

283

But the op is not a _scope_. It's an allocation.

[wsmoses]

Remove scope from name

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (index.hasValue()) {
     return;
   }
 
   regions.push_back(RegionSuccessor(&bodyRegion()));
 }
 
 /// 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;
   for (auto res : op->getResults()) {
     if (auto effect =
             interface.getEffectOnValue<MemoryEffects::Allocate>(res)) {
       if (isa<SideEffects::AutomaticAllocationScopeResource>(
               effect->getResource()))
         return true;
     }
   }
   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;
   for (auto res : op->getResults()) {
     if (auto effect =
             interface.getEffectOnValue<MemoryEffects::Allocate>(res)) {
       if (isa<SideEffects::AutomaticAllocationScopeResource>(
               effect->getResource()))

 struct AllocaScopeInliner : public OpRewritePattern<AllocaScopeOp> {
   using OpRewritePattern<AllocaScopeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AllocaScopeOp op,
                                 PatternRewriter &rewriter) const override {
     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();
     }
 
     // Only apply to if this is this last non-terminator
     // op in the block (lest lifetime be extended) of a one
     // block region
for (auto &r : lastParentWithoutScope->getRegions()) {
                "only one region can contain the op");
         containingRegion = &r;
       }
     }
     assert(containingRegion && "op must be contained in a region");
 
     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
       // lastParentWithoutScope (i.e. where we would hoist to), skip.
       if (llvm::any_of(alloc->getOperands(), [&](Value v) {
             return containingRegion->isAncestor(v.getParentRegion());
           }))
         return WalkResult::skip();

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK:         %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
 // CHECK:         "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
 // CHECK:       }
 // CHECK:       "test.op"() : () -> ()
 // CHECK:       "test.terminator"() : () -> ()
 // 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