[mlir][bufferization] Fix regionOperatesOnMemrefValues (#75016)

`Region::walk([](Block *b) {...})` does not enumerate blocks that are
direct children of the region. These blocks must be checked manually.

Author: matthias-springer

mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

@@ -463,7 +463,7 @@ BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
 }
 
 static bool regionOperatesOnMemrefValues(Region &region) {
-  WalkResult result = region.walk([](Block *block) {
+  auto checkBlock = [](Block *block) {
     if (llvm::any_of(block->getArguments(), isMemref))
       return WalkResult::interrupt();
     for (Operation &op : *block) {
@@ -473,8 +473,18 @@ static bool regionOperatesOnMemrefValues(Region &region) {
         return WalkResult::interrupt();
     }
     return WalkResult::advance();
-  });
-  return result.wasInterrupted();
+  };
+  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

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir

@@ -531,8 +531,8 @@ func.func @noRegionBranchOpInterface() {
 // This is not allowed in buffer deallocation.
 
 func.func @noRegionBranchOpInterface() {
-  // expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
   %0 = "test.bar"() ({
+    // expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
     %1 = "test.bar"() ({
       %2 = "test.get_memref"() : () -> memref<2xi32>
       "test.yield"(%2) : (memref<2xi32>) -> ()
@@ -544,6 +544,21 @@ func.func @noRegionBranchOpInterface() {
 
 // -----
 
+// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface.
+// This is not allowed in buffer deallocation.
+
+func.func @noRegionBranchOpInterface() {
+  // expected-error@+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)
+    "test.yield"(%3) : (i32) -> ()
+  }) : () -> (i32)
+  "test.terminator"() : () -> ()
+}
+
+// -----
+
 func.func @while_two_arg(%arg0: index) {
   %a = memref.alloc(%arg0) : memref<?xf32>
   scf.while (%arg1 = %a, %arg2 = %a) : (memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>) {