[Mlir-commits] [mlir] [mlir][bufferization] Fix ownership computation of unknown ops (PR #70773)

[Matthias Springer]

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

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

>From 3f207024de3dbe6c36cf3a704df8bd8a21edf0b8 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 31 Oct 2023 15:24:27 +0900
Subject: [PATCH] [mlir][bufferization] Fix ownership computation of unknown
 ops

No ownership is assumed for memref results of ops that implement none of the relevant interfaces and have no memref operands. This fixed #68948.
---
 .../OwnershipBasedBufferDeallocation.cpp      | 24 ++++++++------
 .../dealloc-other.mlir                        | 31 +++++++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index b2b77eda92ef210..9459cc43547fafc 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/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));
+    }
   }
 }
 
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir
new file mode 100644
index 000000000000000..be894dbe4c5fc80
--- /dev/null
+++ b/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>
+}