[llvm] [mlir][bufferization] Don't clone on unknown ownership and verify function boundary ABI (PR #66626)

[Martin Erhart via llvm-commits]

================
@@ -132,30 +132,79 @@ void DeallocationState::getLiveMemrefsIn(Block *block,
   memrefs.append(liveMemrefs);
 }
 
-std::pair<Value, Value>
-DeallocationState::getMemrefWithUniqueOwnership(OpBuilder &builder,
-                                                Value memref, Block *block) {
-  auto iter = ownershipMap.find({memref, block});
-  assert(iter != ownershipMap.end() &&
-         "Value must already have been registered in the ownership map");
-
-  Ownership ownership = iter->second;
-  if (ownership.isUnique())
-    return {memref, ownership.getIndicator()};
-
-  // Instead of inserting a clone operation we could also insert a dealloc
-  // operation earlier in the block and use the updated ownerships returned by
-  // the op for the retained values. Alternatively, we could insert code to
-  // check aliasing at runtime and use this information to combine two unique
-  // ownerships more intelligently to not end up with an 'Unknown' ownership in
-  // the first place.
-  auto cloneOp =
-      builder.create<bufferization::CloneOp>(memref.getLoc(), memref);
-  Value condition = buildBoolValue(builder, memref.getLoc(), true);
-  Value newMemref = cloneOp.getResult();
-  updateOwnership(newMemref, condition);
-  memrefsToDeallocatePerBlock[newMemref.getParentBlock()].push_back(newMemref);
-  return {newMemref, condition};
+Value DeallocationState::getMemrefWithUniqueOwnership(
+    const DeallocationOptions &options, OpBuilder &builder, Value memref,
+    Block *block) {
+  // NOTE: * if none of the operands have the same allocated pointer, a new
+  // memref was allocated and thus the operation should have the allocate
+  // side-effect defined on that result value and thus the correct unique
+  // ownership is pre-populated by the ownership pass (unless an interface
+  // implementation is incorrect).
+  //       * if exactly one operand has the same allocated pointer, this retunes
+  //       the ownership of exactly that operand
+  //       * if multiple operands match the allocated pointer of the result, the
+  //       ownership indicators of all of them always have to evaluate to the
+  //       same value because no dealloc operations may be present and because
+  //       of the rules they are passed to nested regions and successor blocks.
+  //       This could be verified at runtime by inserting `cf.assert`
+  //       operations, but would require O(|operands|^2) additional operations
+  //       to check and is thus not implemented yet (would need to insert a
+  //       library function to avoid code-size explosion which would make the
+  //       deallocation pass a module pass)
+  auto ipSave = builder.saveInsertionPoint();
+  SmallVector<Value> worklist;
+  worklist.push_back(memref);
+
+  while (!worklist.empty()) {
----------------
maerhart wrote:

To only materialize the ownership indicators when needed, but maybe it's better to just eagerly materialize them and rely on DCE to get rid of the ones that are not needed. I also just noticed that the interface implementations that materialize custom ownership indicator logic is not queried here. It would probably also allow us to get rid of the `Ownership` datastructure because 'Unknown' ownerships could not occur anymore (I think).

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