[MLIR][MemRef] Nested allocation scope inlining

[wsmoses]

Summary

If a stack allocation is within a nested allocation scope
don't count that as an allocation of the outer allocation scope
that would prevent inlining.

[bondhugula]

@wsmoses Looks like you missed my comments here: https://reviews.llvm.org/D121845 -- before merging it.

[ftynse]

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

317–332

Would it make sense to switch the walk order to preorder (template argument) and check alloc->hasTrait<OpTrait::AutomaticAllocationScope>() + skip instead?

In any case, please do something with the formatting here, e.g., put alloc->getParentWithTrait into a variable.

329–330

Nit: fold nested ifs?

[wsmoses]

Address comments and fix formatting

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

 
 /// Inline an AllocaScopeOp if either the direct parent is an allocation scope
 /// or it contains no allocation.
 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 (alloc == op)
-                return WalkResult::advance();
-              if (isOpItselfPotentialAutomaticAllocation(alloc))
-                return WalkResult::interrupt();
-              return WalkResult::advance();
-            }).wasInterrupted();
-      if (hasPotentialAlloca)
-        return failure();
-    }
+        op->walk<WalkOrder::PreOrder>([&](Operation *alloc) {
+            if (alloc->hasTrait<OpTrait::AutomaticAllocationScope>())
+              return WalkResult::skip();
 
-    // Only apply to if this is this last non-terminator
-    // op in the block (lest lifetime be extended) of a one
-    // block region
+    // If this contains no potential allocation, it is always legal to
+    // inline. Otherwise, consider two conditions:
+    if (hasPotentialAlloca) {
+      // If the parent isn't an allocation scope, or we are not the last
+      // non-terminator op in the parent, we will extend the lifetime.
+      if (!op->getParentOp()->hasTrait<OpTrait::AutomaticAllocationScope>())
+        return failure();
       if (!lastNonTerminatorInRegion(op))
         return failure();
+    }
 
     Block *block = &op.getRegion().front();
     Operation *terminator = block->getTerminator();
     ValueRange results = terminator->getOperands();
     rewriter.mergeBlockBefore(block, op);
     rewriter.replaceOp(op, results);
     rewriter.eraseOp(terminator);
     return success();

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK:         "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
 // CHECK:       }
 // CHECK:       "test.op"() : () -> ()
 // CHECK:       "test.terminator"() : () -> ()
 // CHECK:     }) : () -> ()
 // CHECK:     return
 // CHECK:   }
 
+func @scopeMerge5() {
+  "test.region"() ({
+    memref.alloca_scope {
+      affine.parallel (%arg) = (0) to (64) {
+        %a = memref.alloca(%arg) : memref<?xi64>
+        "test.use"(%a) : (memref<?xi64>) -> ()
+      }
+    }
+    "test.op"() : () -> ()
+    "test.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK:   func @scopeMerge5() {
+// CHECK:     "test.region"() ({
+// CHECK:       affine.parallel (%[[cnt:.+]]) = (0) to (64) {
+// 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"() : () -> ()
   }) : () -> ()