[Mlir-commits] [mlir] [mlir][bufferization] Fix handling of indirect function calls (PR #94896)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

<details>
<summary>Changes</summary>

This commit fixes a crash in the ownership-based buffer deallocation pass when indirectly calling a function via SSA value. Such functions must be conservatively assumed to be public.

Fixes #<!-- -->94780.


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


2 Files Affected:

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


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index a8ec111f8c304..bd5c4d4769216 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -822,10 +822,11 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
 
   // Lookup the function operation and check if it has private visibility. If
   // the function is referenced by SSA value instead of a Symbol, it's assumed
-  // to be always private.
+  // to be public. (And we cannot easily change the type of the SSA value
+  // anyway.)
   Operation *funcOp = op.resolveCallable(state.getSymbolTable());
-  bool isPrivate = true;
-  if (auto symbol = dyn_cast<SymbolOpInterface>(funcOp))
+  bool isPrivate = false;
+  if (auto symbol = dyn_cast_or_null<SymbolOpInterface>(funcOp))
     isPrivate = symbol.isPrivate() && !symbol.isDeclaration();
 
   // If the private-function-dynamic-ownership option is enabled and we are
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 63947150c23b3..a77442dfe40e5 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
@@ -131,3 +131,22 @@ func.func @g(%arg0: 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>
+
+// -----
+
+func.func @func_call_indirect(%m: memref<?xf32>, %f: (memref<?xf32>) -> (memref<?xf32>)) {
+  %0 = func.call_indirect %f(%m) : (memref<?xf32>) -> (memref<?xf32>)
+  return
+}
+
+// CHECK-LABEL: func @func_call_indirect(
+//       CHECK:   %[[true:.*]] = arith.constant true
+//       CHECK:   %[[call:.*]] = call_indirect {{.*}} : (memref<?xf32>) -> memref<?xf32>
+//       CHECK:   %[[base_call:.*]], %{{.*}}, %{{.*}}, %{{.*}} = memref.extract_strided_metadata %[[call]]
+//       CHECK:   bufferization.dealloc (%[[base_call]] : {{.*}}) if (%[[true]])
+
+// CHECK-DYNAMIC-LABEL: func @func_call_indirect(
+//       CHECK-DYNAMIC:   %[[true:.*]] = arith.constant true
+//       CHECK-DYNAMIC:   %[[call:.*]] = call_indirect {{.*}} : (memref<?xf32>) -> memref<?xf32>
+//       CHECK-DYNAMIC:   %[[base_call:.*]], %{{.*}}, %{{.*}}, %{{.*}} = memref.extract_strided_metadata %[[call]]
+//       CHECK-DYNAMIC:   bufferization.dealloc (%[[base_call]] : {{.*}}) if (%[[true]])

``````````

</details>


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