[mlir][bufferization] Fix handling of indirect function calls

[matthias-springer]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

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 --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]])

[maerhart]

Looks good, thanks!