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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

<details>
<summary>Changes</summary>

* 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

cc @<!-- -->Hardcode84 

---
Full diff: https://github.com/llvm/llvm-project/pull/67418.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+10-4) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 94a26f3aff5e019..5ea60e9435c2f52 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);
   }
@@ -802,8 +804,10 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
   // to be always private.
   Operation *funcOp = op.resolveCallable(state.getSymbolTable());
   bool isPrivate = true;
-  if (auto symbol = dyn_cast<SymbolOpInterface>(funcOp))
+  if (auto symbol = dyn_cast<SymbolOpInterface>(funcOp)) {
     isPrivate &= (symbol.getVisibility() == SymbolTable::Visibility::Private);
+    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`
@@ -931,8 +935,10 @@ BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {
 
 bool BufferDeallocation::isFunctionWithoutDynamicOwnership(Operation *op) {
   auto funcOp = dyn_cast<FunctionOpInterface>(op);
-  return funcOp && (!options.privateFuncDynamicOwnership ||
-                    funcOp.getVisibility() != SymbolTable::Visibility::Private);
+  return funcOp &&
+         (!options.privateFuncDynamicOwnership ||
+          funcOp.getVisibility() != SymbolTable::Visibility::Private ||
+          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>

``````````

</details>


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