[mlir][linalg][bufferize] Add `createDealloc` options

[springerm]

Summary

If createDealloc is deactivated (enabled by default), newly allocated buffers are not deallocated anymore. In such a case, the missing deallocations can be inserted by the existing "BufferDeallocation" pass.

This change is needed for unifying core bufferization and Comprehensive Bufferize. Core bufferization has a separate pass for generating deallocations.

Note: In the future, this will evolve towards generating deallocation ops only for buffer allocations that do not escape block boundaries (i.e., that are in destination passing style).

Depends On D116449

[springerm]

update

[nicolasvasilache]

If createDealloc is true (default), newly allocated buffers are deallocated.
Otherwise, the missing deallocations should be inserted by explicitly calling the existing "BufferDeallocation" pass.

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

380

can we drop the default here ? it does not seem to bring value

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

struct BufferizationOptions {
   /// Otherwise, a pass failure is triggered.
   bool allowReturnMemref = false;
 
   /// Specifies whether not bufferizable ops are allowed in the input. If so,
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
+  /// Specifies whether dealloc ops should be generated along with alloc ops. If
+  /// not, new memory allocations will leak.
+  bool createDeallocs = true;
+
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
public:
   /// Note: When reaching an end of the reverse SSA use-def chain, that value
   /// is returned regardless of whether it is a memory write or not.
   Value findLastPrecedingWrite(Value value) const;
 
   /// Creates a memref allocation.
   Optional<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                               ArrayRef<Value> dynShape) const;
 
-  /// Creates an alloc-dealloc pair. This function may perform additional
-  /// optimizations such as buffer allocation hoisting.
-  Value createAllocDeallocPair(OpBuilder &builder, Location loc,
-                               Value shapedValue) const;
+  /// Creates a memref allocation for the given shaped value. This function may
+  /// perform additional optimizations such as buffer allocation hoisting. If
+  /// `createDealloc`, a deallocation op is inserted at the point where the
+  /// allocation goes out of scope.
+  Value createAlloc(OpBuilder &b, Location loc, Value shapedValue,
+                    bool deallocMemref) const;
 
   /// Creates a memref deallocation. The given memref buffer must have been
   /// allocated using `createAlloc`.
   void createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer) const;
 
   /// Creates a memcpy between two given buffers.
   void createMemCpy(OpBuilder &b, Location loc, Value from, Value to) const;
 

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

if (!aliasInfo.isInPlace(result)) {
     assert(
         aliasingOperands.size() == 1 &&
         "ops with multiple aliasing OpOperands cannot bufferize out-of-place");
     Location loc = op->getLoc();
     // Move insertion point right after `operandBuffer`. That is where the
     // allocation should be inserted (in the absence of allocation hoisting).
     setInsertionPointAfter(rewriter, operandBuffer);
     // Allocate the result buffer.
-    Value resultBuffer = createAllocDeallocPair(rewriter, loc, operandBuffer);
+    Value resultBuffer =
+        createAlloc(rewriter, loc, operandBuffer, options.createDeallocs);
     bool skipCopy = false;
     // Do not copy if the last preceding write of `operand` is an op that does
     // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
     // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
     // use-def chain, it returns that value, regardless of whether it is a
     // memory write or not.
     Value lastWrite = findLastPrecedingWrite(operand);
     if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
   // TODO: when packing, allocate a static bounding box which will enable more
   // hoisting.
   if (dynShape.empty())
     moveInsertionPointToAllocationHoistingBarrier(b);
 
   return allocMemRefType;
 }
 
-/// Create an Allocop/DeAllocOp pair, where the AllocOp is after
+/// Create an AllocOp/DeallocOp pair, where the AllocOp is after
 /// `shapedValue.getDefiningOp` (or at the top of the block in case of a
 /// bbArg) and the DeallocOp is at the end of the block.
-Value mlir::linalg::comprehensive_bufferize::BufferizationState::
-    createAllocDeallocPair(OpBuilder &b, Location loc,
-                           Value shapedValue) const {
+Value mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
+    OpBuilder &b, Location loc, Value shapedValue, bool deallocMemref) const {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // 1. Create memory allocation.
   assert(shapedValue.getType().isa<ShapedType>());
   MemRefType memRefType = shapedValue.getType().dyn_cast<MemRefType>();
   SmallVector<Value> dynShape;
   // Note: getAllocationTypeAndShape also sets the insertion point.
   MemRefType allocMemRefType =
       getAllocationTypeAndShape(b, loc, shapedValue, dynShape);
   Optional<Value> allocated = createAlloc(b, loc, allocMemRefType, dynShape);
   // TODO: For now just assert the value is returned. Eventually need to
   // error-propagate.
   assert(allocated && "allocation failed");
   Value casted = allocated.getValue();
   if (memRefType && memRefType != allocMemRefType) {
     casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
   }
 
+  if (deallocMemref) {
     // 2. Create memory deallocation.
     b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
     createDealloc(b, loc, allocated.getValue());
+  }
+
   return casted;
 }
 
 /// Create a memref allocation.
 Optional<Value>
 mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
     OpBuilder &b, Location loc, MemRefType type,
     ArrayRef<Value> dynShape) const {

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp

struct InitTensorOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
-    Value alloc = state.createAllocDeallocPair(rewriter, initTensorOp->getLoc(),
-                                               initTensorOp.result());
+    Value alloc = state.createAlloc(rewriter, initTensorOp->getLoc(),
+                                    initTensorOp.result(),
+                                    state.getOptions().createDeallocs);
     replaceOpWithBufferizedValues(rewriter, op, alloc);
     return success();
   }
 };
 
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
                                                     linalg::TiledLoopOp> {

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     bool inplace = state.isInPlace(extractSliceOp->getResult(0));
     Value alloc;
     if (!inplace)
-      alloc =
-          state.createAllocDeallocPair(rewriter, loc, extractSliceOp.result());
+      alloc = state.createAlloc(rewriter, loc, extractSliceOp.result(),
+                                state.getOptions().createDeallocs);
 
     // Bufferize to subview.
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             dstTensorType.getRank(), srcMemrefType,
             extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
             extractSliceOp.getMixedStrides())
             .cast<MemRefType>();