[mlir][bufferize][NFC] Deallocate all buffers at the end of bufferization

[springerm]

Summary

This makes bufferization more modular. This is in preparation of future refactorings.

Depends On D121361

[springerm]

update

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

 };
 
 /// BufferizationState provides helper functions for performing bufferization
 /// rewrites and handling memref buffers.
 struct BufferizationState {
   BufferizationState(const AnalysisState &analysisState)
       : analysisState(analysisState) {}
 
+  /// Creates a memref allocation with the given type and dynamic extents.
+  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                               ValueRange dynShape);
+
+  /// Creates a memref allocation for the given shaped value. This function may
+  /// perform additional optimizations such as buffer allocation hoisting.
+  // TODO: Allocation hoisting should be a cleanup pass.
+  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, Value shapedValue);
+
   /// Return the buffer (memref) for a given OpOperand (tensor). Allocate
   /// a new buffer and copy over data from the existing buffer if out-of-place
   /// bufferization was decided.
   FailureOr<Value>
   getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
             bool forceInPlace = false,
-            Optional<Operation *> customCopyInsertionPoint = None) const;
+            Optional<Operation *> customCopyInsertionPoint = None);
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const {
     return analysisState.getOptions();
   }
 
   const AnalysisState &getAnalysisState() const { return analysisState; }
 

 /// Return a MemRefType to which the `tensorType` can be bufferized in a
 /// composable fashion. The layout must be the most dynamic possible and
 /// canonicalize away once bufferization is finished.
 BaseMemRefType getMemRefType(TensorType tensorType,
                              const BufferizationOptions &options,
                              MemRefLayoutAttrInterface layout = {},
                              Attribute memorySpace = {});
 
-/// Creates a memref allocation with the given type and dynamic extents.
-FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                             ValueRange dynShape,
-                             const BufferizationOptions &options);
-
-/// Creates a memref allocation with the given type and dynamic extents. If
-/// `createDealloc`, a deallocation op is inserted at the point where the
-/// allocation goes out of scope.
-FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                             ValueRange dynShape, bool deallocMemref,
-                             const BufferizationOptions &options);
-
-/// 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.
-// TODO: Allocation hoisting should be a cleanup pass.
-FailureOr<Value> createAlloc(OpBuilder &b, Location loc, Value shapedValue,
-                             bool deallocMemref,
-                             const BufferizationOptions &options);
-
 /// Creates a memref deallocation. The given memref buffer must have been
 /// allocated using `createAlloc`.
 LogicalResult createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
                             const BufferizationOptions &options);
 
 /// Creates a memcpy between two given buffers.
 LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to,
                            const BufferizationOptions &options);
 
+/// Finalize all buffer allocations, i.e., create alloc ops as specified in the
+/// bufferization options and deallocate all buffers.
+LogicalResult finalizeBuffers(Operation *op,
+                              const BufferizationOptions &options);
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 namespace mlir {
 namespace bufferization {
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

 /// Note: The layout map chosen to bufferize is the most dynamic canonical
 /// strided layout of the proper rank. This ensures compatibility with expected
 /// layouts after transformations. Combinations of memref.cast +
 /// canonicalization are responsible for clean ups.
 // TODO: Extract `options` from `state` and pass as separate argument.
 LogicalResult bufferizeOp(Operation *op, const AnalysisState &analysisState);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
-/// Reuse an existing `BufferizationState`.
-///
-/// Note: This function overload is useful for extending the bufferization.
-LogicalResult bufferizeOp(Operation *op,
-                          BufferizationState &bufferizationState);
-
-/// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Buffers are duplicated and copied before any tensor use that bufferizes to
 /// a memory write.
 ///
 /// Note: This function bufferizes ops without utilizing analysis results. It
 /// can be used to implement partial bufferization passes.
 LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options);
 
 BufferizationOptions getPartialBufferizationOptions();
 
+//===----------------------------------------------------------------------===//
+// Helper functions for extending Bufferization
+//===----------------------------------------------------------------------===//
+
+/// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
+/// Reuse an existing `BufferizationState`.
+///
+/// Note: This function overload is useful for extending the bufferization.
+LogicalResult bufferizeOp(Operation *op,
+                          BufferizationState &bufferizationState);
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

 constexpr const ::llvm::StringLiteral
     bufferization::BufferizableOpInterface::kBufferLayoutAttrName;
 
 /// Attribute name used to mark region arguments that can be bufferized
 /// in-place during linalg comprehensive bufferization.
 constexpr const ::llvm::StringLiteral
     bufferization::BufferizableOpInterface::kInplaceableAttrName;
 
+static const char *kBufferAllocationAttr = "bufferization.allocation";
+
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
 // Default constructor for BufferizationOptions.
 BufferizationOptions::BufferizationOptions() = default;
 
 BufferizableOpInterface
Value mlir::bufferization::lookupBuffer(RewriterBase &rewriter, Value tensor,
   ensureToMemrefOpIsValid(tensor, memrefType);
   return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
                                                     tensor);
 }
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
-FailureOr<Value> BufferizationState::getBuffer(
-    RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace,
-    Optional<Operation *> customCopyInsertionPoint) const {
+FailureOr<Value>
+BufferizationState::getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
+                              bool forceInPlace,
+                              Optional<Operation *> customCopyInsertionPoint) {
   const BufferizationOptions &options = analysisState.getOptions();
   OpBuilder::InsertionGuard guard(rewriter);
   Operation *op = opOperand.getOwner();
   Location loc = op->getLoc();
   Value operand = opOperand.get();
   Value operandBuffer = lookupBuffer(rewriter, operand, options);
 
   if (forceInPlace || analysisState.isInPlace(opOperand))
     return operandBuffer;
 
   // Bufferizing out-of-place: Allocate a new buffer.
   // 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.
-  FailureOr<Value> resultBuffer = createAlloc(rewriter, loc, operandBuffer,
-                                              options.createDeallocs, options);
+  FailureOr<Value> resultBuffer = createAlloc(rewriter, loc, operandBuffer);
   if (failed(resultBuffer))
     return failure();
   // Do not copy if the last preceding writes of `operand` are ops that do
   // 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.
   SetVector<Value> lastWrites = analysisState.findLastPrecedingWrite(operand);
bool AlwaysCopyAnalysisState::areEquivalentBufferizedValues(Value v1,
   // conservative answer is "false".
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//
 
+/// Create a memref allocation with the given type and dynamic extents.
+static FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                                    ValueRange dynShape,
+                                    const BufferizationOptions &options) {
+  if (options.allocationFn)
+    return (*options.allocationFn)(b, loc, type, dynShape,
+                                   options.bufferAlignment);
+
+  // Default bufferallocation via AllocOp.
+  Value allocated = b.create<memref::AllocOp>(
+      loc, type, dynShape, b.getI64IntegerAttr(options.bufferAlignment));
+  return allocated;
+}
+
+/// Creates a memref deallocation. The given memref buffer must have been
+/// allocated using `createAlloc`.
+LogicalResult
+bufferization::createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
+                             const BufferizationOptions &options) {
+  if (options.deallocationFn)
+    return (*options.deallocationFn)(b, loc, allocatedBuffer);
+
+  // Default buffer deallocation via DeallocOp.
+  b.create<memref::DeallocOp>(loc, allocatedBuffer);
+  return success();
+}
+
 /// Move the insertion point of the given builder to the beginning of a
 /// surrounding block as much as possible, while not crossing any allocation
 /// hoisting barriers.
 static void moveInsertionPointToAllocationHoistingBarrier(OpBuilder &b) {
   Operation *op = b.getInsertionBlock()->getParentOp();
   while (op) {
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       if (bufferizableOp.isAllocationHoistingBarrier())
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
-/// `shapedValue.getDefiningOp` (or at the top of the block in case of a
-/// bbArg) and the DeallocOp is at the end of the block.
-FailureOr<Value>
-bufferization::createAlloc(OpBuilder &b, Location loc, Value shapedValue,
-                           bool deallocMemref,
-                           const BufferizationOptions &options) {
+static Value createBufferAllocation(OpBuilder &b, Location loc, MemRefType type,
+                                    ValueRange dynShape) {
+  auto allocaOp = b.create<memref::AllocaOp>(loc, type, dynShape);
+  allocaOp->setAttr(kBufferAllocationAttr, b.getUnitAttr());
+  return allocaOp.getResult();
+}
+
+/// Create an allocation after `shapedValue.getDefiningOp` (or at the top of the
+/// block in case of a bbArg).
+FailureOr<Value> BufferizationState::createAlloc(OpBuilder &b, Location loc,
+                                                 Value shapedValue) {
   // 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);
-  FailureOr<Value> allocated =
-      createAlloc(b, loc, allocMemRefType, dynShape, options);
-  if (failed(allocated))
-    return failure();
-  Value casted = allocated.getValue();
+  Value alloc = createBufferAllocation(b, loc, allocMemRefType, dynShape);
   if (memRefType && memRefType != allocMemRefType) {
-    assert(memref::CastOp::areCastCompatible(allocated.getValue().getType(),
-                                             memRefType) &&
+    assert(memref::CastOp::areCastCompatible(alloc.getType(), memRefType) &&
            "createAlloc: cast incompatible");
-    casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
-  }
-
-  if (deallocMemref) {
-    // 2. Create memory deallocation.
-    b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
-    if (failed(createDealloc(b, loc, allocated.getValue(), options)))
-      return failure();
-  }
-
-  return casted;
-}
-
-/// Create a memref allocation with the given type and dynamic extents.
-FailureOr<Value>
-bufferization::createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                           ValueRange dynShape,
-                           const BufferizationOptions &options) {
-  if (options.allocationFn)
-    return (*options.allocationFn)(b, loc, type, dynShape,
-                                   options.bufferAlignment);
-
-  // Default bufferallocation via AllocOp.
-  Value allocated = b.create<memref::AllocOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(options.bufferAlignment));
-  return allocated;
+    alloc = b.create<memref::CastOp>(loc, memRefType, alloc);
   }
-
-/// Create a memref allocation with the given type and dynamic extents. May also
-/// deallocate the memref again.
-FailureOr<Value>
-bufferization::createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                           ValueRange dynShape, bool deallocMemref,
-                           const BufferizationOptions &options) {
-  OpBuilder::InsertionGuard g(b);
-
-  FailureOr<Value> alloc = createAlloc(b, loc, type, dynShape, options);
-  if (failed(alloc))
-    return failure();
-
-  if (deallocMemref) {
-    // Dealloc at the end of the block.
-    b.setInsertionPoint(alloc.getValue().getParentBlock()->getTerminator());
-    if (failed(createDealloc(b, loc, *alloc, options)))
-      return failure();
-  }
-
   return alloc;
 }
 
-/// Create a memref deallocation.
-LogicalResult
-bufferization::createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
-                             const BufferizationOptions &options) {
-  if (options.deallocationFn)
-    return (*options.deallocationFn)(b, loc, allocatedBuffer);
-
-  // Default buffer deallocation via DeallocOp.
-  b.create<memref::DeallocOp>(loc, allocatedBuffer);
-  return success();
+/// Create a memref allocation with the given type and dynamic extents.
+FailureOr<Value> BufferizationState::createAlloc(OpBuilder &b, Location loc,
+                                                 MemRefType type,
+                                                 ValueRange dynShape) {
+  return createBufferAllocation(b, loc, type, dynShape);
 }
 
 /// Create a memory copy between two memref buffers.
 LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
                                           Value from, Value to,
                                           const BufferizationOptions &options) {
   if (options.memCpyFn)
     return (*options.memCpyFn)(b, loc, from, to);
 
   b.create<memref::CopyOp>(loc, from, to);
   return success();
 }
 
+LogicalResult
+bufferization::finalizeBuffers(Operation *op,
+                               const BufferizationOptions &options) {
+  IRRewriter rewriter(op->getContext());
+
+  // 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();
+
+    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);
+
+    // Stop here if deallocations are deactivated.
+    if (!options.createDeallocs)
+      return WalkResult::advance();
+
+    rewriter.setInsertionPoint(block->getTerminator());
+    if (failed(createDealloc(rewriter, alloc->getLoc(), *alloc, options)))
+      return WalkResult::interrupt();
+
+    return WalkResult::advance();
+  });
+
+  return success(!status.wasInterrupted());
+}
+
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 bool bufferization::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;

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

checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
   }
 
   return success();
 }
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const AnalysisState &analysisState) {
   BufferizationState bufferizationState(analysisState);
-  return bufferizeOp(op, bufferizationState);
+  if (failed(bufferizeOp(op, bufferizationState)))
+    return failure();
+  if (failed(finalizeBuffers(op, analysisState.getOptions())))
+    return failure();
+  return success();
 }
 
 LogicalResult
 bufferization::bufferizeOp(Operation *op,
                            BufferizationState &bufferizationState) {
   // Bufferize the op and its nested ops.
   RewritePatternSet patterns(op->getContext());
   patterns.add<BufferizationPattern>(patterns.getContext(), bufferizationState);
bufferization::bufferizeOp(Operation *op,
   // TODO: Perform a preorder walk instead of the greedy pattern rewriter. This
   // would be more efficient because every bufferization pattern is guaranteed
   // to apply only a single time (otherwise, an assertion would be triggered).
   // However, there are restrictions wrt. erasing ops during a preorder walk,
   // which would likely require a larger refactoring.
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns), config)))
     return failure();
 
-  return checkBufferizationResult(op, bufferizationState.getOptions());
+  if (failed(checkBufferizationResult(op, bufferizationState.getOptions())))
+    return failure();
+
+  return success();
 }
 
 namespace {
 /// This a "no analysis, always copy" AnalysisState. In the absence of an
 /// analysis, a buffer must be copied each time it is written to. Therefore, all
 /// OpOperands that bufferize to a memory write must bufferize out-of-place.
 class AlwaysCopyAnalysisState : public AnalysisState {
 public:

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

if (!options.allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }
 
+  // Finalize all buffers.
+  if (failed(finalizeBuffers(moduleOp, options)))
+    return failure();
+
   // Perform a post-processing pass of layout modification at function boundary
   // according to the kBufferLayoutAttrName.
   layoutPostProcessing(moduleOp);
 
   // Post-pass cleanup of inplaceable and buffer_layout attributes.
   moduleOp.walk([&](FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationFuncArguments(bbArg);
   });
 
   return success();
 }

mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp

struct InitTensorOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
-    FailureOr<Value> alloc =
-        createAlloc(rewriter, initTensorOp->getLoc(), initTensorOp.result(),
-                    state.getOptions().createDeallocs, state.getOptions());
+    FailureOr<Value> alloc = state.createAlloc(rewriter, initTensorOp->getLoc(),
+                                               initTensorOp.result());
     if (failed(alloc))
       return failure();
     replaceOpWithBufferizedValues(rewriter, op, *alloc);
     return success();
   }
 };
 
 /// Helper structure that iterates over all LinalgOps in `OpTys` and registers

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     bool inplace =
         state.getAnalysisState().isInPlace(extractSliceOp->getOpOperand(0));
     Value alloc;
     if (!inplace) {
       FailureOr<Value> allocOrFailure =
-          createAlloc(rewriter, loc, extractSliceOp.result(),
-                      state.getOptions().createDeallocs, state.getOptions());
+          state.createAlloc(rewriter, loc, extractSliceOp.result());
       if (failed(allocOrFailure))
         return failure();
       alloc = *allocOrFailure;
     }
 
     // Expand offsets, sizes and strides to the full rank to handle the
     // rank-reducing case.
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto fromElementsOp = cast<tensor::FromElementsOp>(op);
 
     // Allocate a buffer for the result.
     Location loc = op->getLoc();
     auto tensorType = fromElementsOp.getType().cast<RankedTensorType>();
     auto shape = tensorType.getShape();
     MemRefType resultType = getContiguousMemRefType(tensorType);
     FailureOr<Value> maybeBuffer =
-        createAlloc(rewriter, loc, resultType, {},
-                    /*deallocMemref=*/state.getOptions().createDeallocs,
-                    state.getOptions());
+        state.createAlloc(rewriter, loc, resultType, {});
     if (failed(maybeBuffer))
       return failure();
     Value buffer = *maybeBuffer;
 
     // Case: tensor<0xelem_type>.
     if (fromElementsOp.elements().empty()) {
       replaceOpWithBufferizedValues(rewriter, op, buffer);
       return success();
struct GenerateOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto generateOp = cast<tensor::GenerateOp>(op);
 
     // Allocate memory.
     Location loc = op->getLoc();
     MemRefType memrefType =
         getContiguousMemRefType(generateOp.getType().cast<RankedTensorType>());
-    FailureOr<Value> maybeResult =
-        createAlloc(rewriter, loc, memrefType, generateOp.dynamicExtents(),
-                    /*deallocMemref=*/state.getOptions().createDeallocs,
-                    state.getOptions());
+    FailureOr<Value> maybeResult = state.createAlloc(
+        rewriter, loc, memrefType, generateOp.dynamicExtents());
     if (failed(maybeResult))
       return failure();
     Value result = *maybeResult;
 
     // Collect loop bounds.
     int64_t rank = memrefType.getRank();
     Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);