[Mlir-commits] [mlir] [mlir][bufferization] Buffer deallocation: skip ops that do not operate on buffers (PR #75126)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/75126

Skip ops that do not operate on buffers during ownership-based buffer deallocation. Such ops can be ignored during buffer deallocation. (Except for terminators.)

>From 9da3a82099cc53aa278ed40881de124e22b3e5a4 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 12 Dec 2023 10:48:28 +0900
Subject: [PATCH] [mlir][bufferization] Buffer deallocation: skip ops that do
 not operate on buffers

Skip ops that do not operate on buffers during ownership-based buffer deallocation. Such ops can be ignored during buffer deallocation. (Except for terminators.)
---
 .../OwnershipBasedBufferDeallocation.cpp      | 48 ++++++++-----------
 .../dealloc-region-branchop-interface.mlir    |  3 +-
 2 files changed, 21 insertions(+), 30 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 38ffae68a43de..f9ec6564995c8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -48,6 +48,19 @@ static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) {
 
 static bool isMemref(Value v) { return v.getType().isa<BaseMemRefType>(); }
 
+/// Return "true" if the given op has buffer semantics. I.e., it has buffer
+/// operands, buffer results and/or buffer region entry block arguments.
+static bool hasBufferSemantics(Operation *op) {
+  if (llvm::any_of(op->getOperands(), isMemref) ||
+      llvm::any_of(op->getResults(), isMemref))
+    return true;
+  for (Region &region : op->getRegions())
+    if (!region.empty())
+      if (llvm::any_of(region.front().getArguments(), isMemref))
+        return true;
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Backedges analysis
 //===----------------------------------------------------------------------===//
@@ -462,31 +475,6 @@ BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
   return state.getMemrefWithUniqueOwnership(builder, memref, block);
 }
 
-static bool regionOperatesOnMemrefValues(Region &region) {
-  auto checkBlock = [](Block *block) {
-    if (llvm::any_of(block->getArguments(), isMemref))
-      return WalkResult::interrupt();
-    for (Operation &op : *block) {
-      if (llvm::any_of(op.getOperands(), isMemref))
-        return WalkResult::interrupt();
-      if (llvm::any_of(op.getResults(), isMemref))
-        return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  };
-  WalkResult result = region.walk(checkBlock);
-  if (result.wasInterrupted())
-    return true;
-
-  // Note: Block::walk/Region::walk visits only blocks that are nested under
-  // nested operations, but not direct children.
-  for (Block &block : region)
-    if (checkBlock(&block).wasInterrupted())
-      return true;
-
-  return false;
-}
-
 LogicalResult
 BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {
   // (1) Ensure that there are supported loops only (no explicit control flow
@@ -512,9 +500,8 @@ LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {
   size_t size = regions.size();
   if (((size == 1 && !op->getResults().empty()) || size > 1) &&
       !dyn_cast<RegionBranchOpInterface>(op)) {
-    if (llvm::any_of(regions, regionOperatesOnMemrefValues))
-      return op->emitError("All operations with attached regions need to "
-                           "implement the RegionBranchOpInterface.");
+    return op->emitError("All operations with attached regions need to "
+                         "implement the RegionBranchOpInterface.");
   }
 
   // (2) The pass does not work properly when deallocations are already present.
@@ -648,6 +635,11 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
   // For each operation in the block, handle the interfaces that affect aliasing
   // and ownership of memrefs.
   for (Operation &op : llvm::make_early_inc_range(*block)) {
+    // Skip ops that do not operate on buffers and are not terminators.
+    // (bufferization.dealloc ops are inserted in front of terminators, so
+    // terminators cannot be skipped.)
+    if (!op.hasTrait<OpTrait::IsTerminator>() && !hasBufferSemantics(&op))
+      continue;
     FailureOr<Operation *> result = handleAllInterfaces(&op);
     if (failed(result))
       return failure();
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
index 1a8a930bc9002..f44b6d74051e2 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
@@ -545,10 +545,9 @@ func.func @noRegionBranchOpInterface() {
 // -----
 
 // Test Case: The op "test.bar" does not implement the RegionBranchOpInterface.
-// This is not allowed in buffer deallocation.
+// But it also does not operate on buffers, so we don't care.
 
 func.func @noRegionBranchOpInterface() {
-  // expected-error at +1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
   %0 = "test.bar"() ({
     %2 = "test.get_memref"() : () -> memref<2xi32>
     %3 = "test.foo"(%2) : (memref<2xi32>) -> (i32)