[Mlir-commits] [mlir] [mlir][bufferization] OwnershipBasedBufferDeallocation fixes (PR #67418)

[Martin Erhart]

https://github.com/maerhart updated https://github.com/llvm/llvm-project/pull/67418

>From 169de4efe4feae032842bb0fe87690234d8e41d1 Mon Sep 17 00:00:00 2001
From: Martin Erhart <merhart at google.com>
Date: Tue, 26 Sep 2023 12:08:04 +0000
Subject: [PATCH] [mlir][bufferization] OwnershipBasedBufferDeallocation fixes

* Properly handle the case where an op is deleted and thus no other interfaces should be processed anymore.
* Don't add ownership indicator arguments and results to function declarations
---
 .../OwnershipBasedBufferDeallocation.cpp        |  8 +++++---
 .../dealloc-callop-interface.mlir               | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 94a26f3aff5e019..f698ef5f9ba08c4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -193,7 +193,7 @@ class BufferDeallocation {
       next = *result;
     }
     if (!next)
-      return nullptr;
+      return FailureOr<Operation *>(nullptr);
     return handleOp<InterfacesU...>(next);
   }
 
@@ -640,6 +640,8 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
     FailureOr<Operation *> result = handleAllInterfaces(&op);
     if (failed(result))
       return failure();
+    if (!*result)
+      continue;
 
     populateRemainingOwnerships(*result);
   }
@@ -803,7 +805,7 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
   Operation *funcOp = op.resolveCallable(state.getSymbolTable());
   bool isPrivate = true;
   if (auto symbol = dyn_cast<SymbolOpInterface>(funcOp))
-    isPrivate &= (symbol.getVisibility() == SymbolTable::Visibility::Private);
+    isPrivate = symbol.isPrivate() && !symbol.isDeclaration();
 
   // If the private-function-dynamic-ownership option is enabled and we are
   // calling a private function, we need to add an additional `i1`
@@ -932,7 +934,7 @@ BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {
 bool BufferDeallocation::isFunctionWithoutDynamicOwnership(Operation *op) {
   auto funcOp = dyn_cast<FunctionOpInterface>(op);
   return funcOp && (!options.privateFuncDynamicOwnership ||
-                    funcOp.getVisibility() != SymbolTable::Visibility::Private);
+                    !funcOp.isPrivate() || funcOp.isExternal());
 }
 
 void BufferDeallocation::populateRemainingOwnerships(Operation *op) {
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
index 389840c9cf77c9a..02bf2d10e9e3f56 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
@@ -114,3 +114,20 @@ func.func @function_call_requries_merged_ownership_mid_block(%arg0: i1) {
 // two allocations was selected, either by checking aliasing of the result at
 // runtime or by extracting the select condition using an OpInterface or by
 // hardcoding the select op
+
+// -----
+
+func.func private @f(memref<f32>) -> memref<f32>
+
+func.func @g(%arg0: memref<f32>) -> memref<f32> {
+  %0 = call @f(%arg0) : (memref<f32>) -> memref<f32>
+  return %0 : memref<f32>
+}
+
+// CHECK-LABEL: func private @f
+//  CHECK-SAME: (memref<f32>) -> memref<f32>
+//       CHECK: call @f({{.*}}) : (memref<f32>) -> memref<f32>
+
+// CHECK-DYNAMIC-LABEL: func private @f
+//  CHECK-DYNAMIC-SAME: (memref<f32>) -> memref<f32>
+//       CHECK-DYNAMIC: call @f({{.*}}) : (memref<f32>) -> memref<f32>