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

[Matthias Springer]

================
@@ -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()) {
----------------
matthias-springer wrote:

That sounds like an interesting idea. (I think that could be a separate NFC pull request?) Can we run into issues where certain unneeded ownership indicators cannot fold away because canonicalizers (for loops, unstructured control flow, etc.) are not efficient enough and the resulting IR is less efficient?

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