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

[llvmlistbot at llvm.org]

Author: Martin Erhart
Date: 2023-09-27T09:16:43+02:00
New Revision: 77813b088eb714ecb888c6804b265523fa254eb7

URL: https://github.com/llvm/llvm-project/commit/77813b088eb714ecb888c6804b265523fa254eb7
DIFF: https://github.com/llvm/llvm-project/commit/77813b088eb714ecb888c6804b265523fa254eb7.diff

LOG: [mlir][bufferization] OwnershipBasedBufferDeallocation fixes (#67418)

* 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

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
    mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir

Removed: 
    


################################################################################
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>