[mlir][bufferize] Do not run the buffer deallocation pass if no allocs escape block boundaries

[springerm]

Summary

This fixes a bufferization issue with ops that are not supported by the buffer deallocation pass when allow-return-allocs=0.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to,
                            const BufferizationOptions &options);
 
 /// Try to hoist all new buffer allocations until the next hoisting barrier.
 LogicalResult hoistBufferAllocations(Operation *op,
                                      const BufferizationOptions &options);
 
 /// Create alloc/dealloc ops as specified in the bufferization options. If
 /// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
-/// buffer deallocation can be created.
+/// buffer deallocation can be created. `changed` is set to `true` if the IR was
+/// modified.
 LogicalResult createAllocDeallocOps(Operation *op,
                                     const BufferizationOptions &options,
-                                    bool onlyLeakingAllocs = false);
+                                    bool onlyLeakingAllocs = false,
+                                    bool *changed = nullptr);
+
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
 
   b.create<memref::CopyOp>(loc, from, to);
   return success();
 }
 
 LogicalResult
 bufferization::createAllocDeallocOps(Operation *op,
                                      const BufferizationOptions &options,
-                                     bool onlyLeakingAllocs) {
+                                     bool onlyLeakingAllocs, bool *changed) {
   IRRewriter rewriter(op->getContext());
+  if (changed)
+    *changed = false;
 
   // Bufferization creates memref.alloca ops. After bufferization, these must be
   // rewritten to alloc/dealloc ops as specified in the bufferization options.
   WalkResult status = op->walk([&](memref::AllocaOp allocaOp) {
     // Ignore memref.alloca ops that were not created by the bufferization.
     if (!allocaOp->hasAttr(kBufferAllocationAttr))
       return WalkResult::skip();
     // If `onlyLeakingAllocs`, process only ops that are marked as
     // "skip dealloc".
     bool skipDealloc = allocaOp->hasAttr(kSkipDeallocAttr);
     if (onlyLeakingAllocs && !skipDealloc)
       return WalkResult::skip();
 
     // Create alloc.
     Block *block = allocaOp->getBlock();
     rewriter.setInsertionPoint(allocaOp);
     FailureOr<Value> alloc =
         createAlloc(rewriter, allocaOp->getLoc(), allocaOp.getType(),
                     allocaOp.dynamicSizes(), options);
     if (failed(alloc))
       return WalkResult::interrupt();
     rewriter.replaceOp(allocaOp, *alloc);
+    if (changed)
+      *changed = true;
 
     // Stop here if the buffer should not be deallocated.
     if (skipDealloc)
       return WalkResult::advance();
 
     // Create dealloc.
     rewriter.setInsertionPoint(block->getTerminator());
     if (failed(createDealloc(rewriter, alloc->getLoc(), *alloc, options)))

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

checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
   }
 
   return success();
 }
 
 LogicalResult
 bufferization::finalizeBuffers(Operation *op,
                                const BufferizationOptions &options) {
+  // Hoist buffers.
   if (failed(hoistBufferAllocations(op, options)))
     return failure();
-  if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true)))
+
+  // Deallocate buffers that escape block boundaries ("leaking buffers") with
+  // the buffer deallocation pass.
+  bool hasLeakingAlloc = false;
+  if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true,
+                                   &hasLeakingAlloc)))
     return failure();
-  if (options.createDeallocs && failed(deallocateBuffers(op)))
+  if (options.createDeallocs && hasLeakingAlloc &&
+      failed(deallocateBuffers(op)))
     return failure();
+
+  // Deallocate all remaining buffers at the end of the block.
   if (failed(createAllocDeallocOps(op, options)))
     return failure();
 
   return success();
 }
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const AnalysisState &analysisState) {