[mlir][bufferization] Fix ownership computation of unknown ops (llvm#70773)

No ownership is assumed for memref results of ops that implement none of the relevant interfaces and have no memref operands. This fixes llvm#68948.

Author: matthias-springer

mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

@@ -973,16 +973,9 @@ void BufferDeallocation::populateRemainingOwnerships(Operation *op) {
  if (!state.getOwnership(res, op->getBlock()).isUninitialized())
  continue;
 
- // Don't take ownership of a returned memref if no allocate side-effect is
- // present, relevant for memref.get_global, for example.
- if (op->getNumOperands() == 0) {
- OpBuilder builder(op);
- state.updateOwnership(res, buildBoolValue(builder, op->getLoc(), false));
- continue;
- }
-
- // Assume the result may alias with any operand and thus combine all their
- // ownerships.
+ // The op does not allocate memory, otherwise, it would have been assigned
+ // an ownership during `handleInterface`. Assume the result may alias with
+ // any memref operand and thus combine all their ownerships.
  for (auto operand : op->getOperands()) {
  if (!isMemref(operand))
  continue;
@@ -991,6 +984,17 @@ void BufferDeallocation::populateRemainingOwnerships(Operation *op) {
  res, state.getOwnership(operand, operand.getParentBlock()),
  op->getBlock());
  }
+
+ // If the ownership value is still uninitialized (e.g., because the op has
+ // no memref operands), assume that no ownership is taken. E.g., this is the
+ // case for "memref.get_global".
+ //
+ // Note: This can lead to memory leaks if memory side effects are not
+ // properly specified on the op.
+ if (state.getOwnership(res, op->getBlock()).isUninitialized()) {
+ OpBuilder builder(op);
+ state.updateOwnership(res, buildBoolValue(builder, op->getLoc(), false));
+ }
  }
 }
 

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir

@@ -0,0 +1,31 @@
+// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation \
+// RUN: --buffer-deallocation-simplification -canonicalize -split-input-file %s | FileCheck %s
+
+// No ownership is assumed for ops that do not implement any interface and have
+// no memref operands.
+
+// CHECK-LABEL: func private @no_interface_no_operands(
+// CHECK-NEXT: %[[m:.*]] = bufferization.to_memref
+// CHECK-NEXT: %[[clone:.*]] = bufferization.clone %[[m]]
+// CHECK-NEXT: return %[[clone]]
+func.func private @no_interface_no_operands(%t : tensor<?x?x?xf16>) -> memref<?x?x?xf16> {
+ %0 = bufferization.to_memref %t : memref<?x?x?xf16>
+ return %0 : memref<?x?x?xf16>
+}
+
+// -----
+
+// If an op does not implement any interface but has memref operands, the
+// ownership of the memref results is computed from the operands.
+
+// CHECK-LABEL: func private @no_interface(
+// CHECK: %[[true:.*]] = arith.constant true
+// CHECK: %[[alloc:.*]] = memref.alloc
+// CHECK: %[[foo:.*]] = "test.foo"(%[[alloc]])
+// CHECK: %[[r:.*]] = bufferization.dealloc (%[[alloc]] : {{.*}}) if (%[[true]]) retain (%[[foo]] : {{.*}})
+// CHECK: return %[[foo]]
+func.func private @no_interface() -> memref<5xf32> {
+ %0 = memref.alloc() : memref<5xf32>
+ %1 = "test.foo"(%0) : (memref<5xf32>) -> (memref<5xf32>)
+ return %1 : memref<5xf32>
+}