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

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/94896

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.


>From efb1a669497bf9f87e05a98a91e6d1e03ba3518c Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Sun, 9 Jun 2024 11:58:04 +0200
Subject: [PATCH] [mlir][bufferization] Fix handling of indirect function calls

This commit fixes a crash in the ownership-based buffer dellocation pass when indirectly calling a function via SSA value. Such functions must be conservatively assumed to be public.
---
 .../OwnershipBasedBufferDeallocation.cpp      |  7 ++++---
 .../dealloc-callop-interface.mlir             | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

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