[mlir][linalg][bufferize][NFC] Change allocationFn return type to FailureOr<Value>

[springerm]

Summary

In addition, all functions that call allocationFn now return FailureOr<Value>. This resolves a few TODOs in the code base.

Depends On D116451

[springerm]

update

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

 class BufferizationState;
 struct PostAnalysisStep;
 
 /// Callback functions that are used to allocate/deallocate/copy memory buffers.
 /// Comprehensive Bufferize provides default implementations of these functions.
 // TODO: Could be replaced with a "bufferization strategy" object with virtual
 // functions in the future.
 struct AllocationCallbacks {
-  using AllocationFn = std::function<Optional<Value>(
+  using AllocationFn = std::function<FailureOr<Value>(
       OpBuilder &, Location, MemRefType, ArrayRef<Value>)>;
   using DeallocationFn = std::function<void(OpBuilder &, Location, Value)>;
   using MemCpyFn = std::function<void(OpBuilder &, Location, Value, Value)>;
 
   AllocationCallbacks(AllocationFn allocFn, DeallocationFn deallocFn,
                       MemCpyFn copyFn)
       : allocationFn(allocFn), deallocationFn(deallocFn), memCpyFn(copyFn) {}
 
public:
   /// Furthermore, BlockArguments are also assumed to be writes. There is no
   /// analysis across block boundaries.
   ///
   /// 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,
+  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                                ArrayRef<Value> dynShape) 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,
+  FailureOr<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

 using namespace linalg::comprehensive_bufferize;
 
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
 /// Default allocation function that is used by the comprehensive bufferization
 /// pass. The default currently creates a ranked memref using `memref.alloc`.
-static Optional<Value> defaultAllocationFn(OpBuilder &b, Location loc,
+static FailureOr<Value> defaultAllocationFn(OpBuilder &b, Location loc,
                                             MemRefType type,
                                             ArrayRef<Value> dynShape) {
   Value allocated = b.create<memref::AllocOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 /// Default deallocation function that is used by the comprehensive
 /// bufferization pass. It expects to recieve back the value called from the
 /// `defaultAllocationFn`.
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 =
+    FailureOr<Value> resultBuffer =
         createAlloc(rewriter, loc, operandBuffer, options.createDeallocs);
+    if (failed(resultBuffer))
+      return failure();
     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))
       if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(), *this))
         skipCopy = true;
     // Do not copy if the copied data is never read.
     if (!isValueRead(result))
       skipCopy = true;
     // Do not copy if this op does not read the data, but writes it.
     if (bufferizesToMemoryWrite(*opOperand) &&
         !bufferizesToMemoryRead(*opOperand))
       skipCopy = true;
     if (!skipCopy) {
       // The copy happens right before the op that is bufferized.
       rewriter.setInsertionPoint(op);
-      createMemCpy(rewriter, loc, operandBuffer, resultBuffer);
+      createMemCpy(rewriter, loc, operandBuffer, *resultBuffer);
     }
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
 
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.
-Value mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
+FailureOr<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");
+  FailureOr<Value> allocated = createAlloc(b, loc, allocMemRefType, dynShape);
+  if (failed(allocated))
+    return failure();
   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>
+FailureOr<Value>
 mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
     OpBuilder &b, Location loc, MemRefType type,
     ArrayRef<Value> dynShape) const {
   return options.allocationFns->allocationFn(b, loc, type, dynShape);
 }
 
 /// Create a memref deallocation.
 void mlir::linalg::comprehensive_bufferize::BufferizationState::createDealloc(

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

static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
   }
 
   // New output operands for the cloned op.
   SmallVector<Value> newOutputBuffers;
   for (OpOperand *opOperand : op.getOutputOperands()) {
     OpResult opResult = op.getTiedOpResult(opOperand);
     assert(opResult && "could not find correspond OpResult");
     FailureOr<Value> resultBuffer = state.getResultBuffer(rewriter, opResult);
+    if (failed(resultBuffer))
+      return failure();
     newOutputBuffers.push_back(*resultBuffer);
   }
 
   // Merge input/output operands.
   SmallVector<Value> newOperands = newInputBuffers;
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
 
   // Set insertion point now that potential alloc/dealloc are introduced.
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.createAlloc(rewriter, initTensorOp->getLoc(),
-                                    initTensorOp.result(),
-                                    state.getOptions().createDeallocs);
-    replaceOpWithBufferizedValues(rewriter, op, alloc);
+    FailureOr<Value> alloc = state.createAlloc(
+        rewriter, initTensorOp->getLoc(), initTensorOp.result(),
+    if (failed(alloc))
+      return failure();
+    replaceOpWithBufferizedValues(rewriter, op, *alloc);
     return success();
   }
 };
 
 /// Bufferization of linalg.tiled_loop. Replace with a new linalg.tiled_loop
 /// that operates entirely on memrefs.
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
for (Value value : tiledLoopOp.inputs()) {
         newInputs.push_back(value);
       }
     }
     int nextResultNum = 0;
     for (Value value : tiledLoopOp.outputs()) {
       if (value.getType().isa<TensorType>()) {
         FailureOr<Value> buffer = state.getResultBuffer(
             rewriter, tiledLoopOp->getResult(nextResultNum++));
+        if (failed(buffer))
+          return failure();
         newOutputs.push_back(*buffer);
         newResults.push_back(*buffer);
       } else {
         newOutputs.push_back(value);
       }
     }
 
     // Create new TiledLoopOp.

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

auto convert = [&](ValueRange values,
         size_t idx = it.index();
         Value val = it.value();
         result.push_back(indices.contains(idx) ? func(val, idx) : val);
       }
       return result;
     };
 
     // Construct a new scf.for op with memref instead of tensor values.
+    bool resultBufferFailure = false;
     SmallVector<Value> initArgs =
         convert(forOp.getInitArgs(), [&](Value val, int64_t index) {
-          return *state.getResultBuffer(rewriter, forOp->getOpResult(index));
+          FailureOr<Value> resultBuffer =
+              state.getResultBuffer(rewriter, forOp->getOpResult(index));
+          if (failed(resultBuffer)) {
+            resultBufferFailure = true;
+            return Value();
+          }
+          return *resultBuffer;
         });
+    if (resultBufferFailure)
+      return failure();
     auto newForOp = rewriter.create<scf::ForOp>(
         forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
         forOp.getStep(), initArgs);
     Block *loopBody = &newForOp.getLoopBody().front();
 
     // Set up new iter_args. The loop body uses tensors, so wrap the (memref)
     // iter_args of the new loop in ToTensorOps.
     rewriter.setInsertionPointToStart(loopBody);

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

struct CastOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
     FailureOr<Value> resultBuffer =
         state.getResultBuffer(rewriter, castOp->getResult(0));
+    if (failed(resultBuffer))
+      return failure();
     auto sourceMemRefType = resultBuffer->getType().cast<BaseMemRefType>();
     Attribute memorySpace = sourceMemRefType.getMemorySpace();
     TensorType resultTensorType =
         castOp.getResult().getType().cast<TensorType>();
     MemRefLayoutAttrInterface layout;
 
     if (auto rankedMemRefType = sourceMemRefType.dyn_cast<MemRefType>())
       if (resultTensorType.isa<RankedTensorType>())
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Value srcMemref = state.lookupBuffer(rewriter, extractSliceOp.source());
     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.createAlloc(rewriter, loc, extractSliceOp.result(),
-                                state.getOptions().createDeallocs);
+    if (!inplace) {
+      FailureOr<Value> allocOrFailure =
+          state.createAlloc(rewriter, loc, extractSliceOp.result(),
+      if (failed(allocOrFailure))
+        return failure();
+      alloc = *allocOrFailure;
+    }
 
     // Bufferize to subview.
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             dstTensorType.getRank(), srcMemrefType,
             extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
             extractSliceOp.getMixedStrides())
             .cast<MemRefType>();
getAliasingOpOperand(Operation *op, OpResult opResult,
     return {&op->getOpOperand(1) /*dest*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     FailureOr<Value> destMemref =
         state.getResultBuffer(rewriter, insertOp->getOpResult(0));
+    if (failed(destMemref))
+      return failure();
     rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.scalar(),
                                      *destMemref, insertOp.indices());
     replaceOpWithBufferizedValues(rewriter, op, *destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // catastrophically bad scheduling decision.
     // TODO: be very loud about it or even consider failing the pass.
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
     Location loc = insertSliceOp.getLoc();
 
     // When bufferizing out-of-place, `getResultBuffer` allocates.
     FailureOr<Value> dstMemref =
         state.getResultBuffer(rewriter, insertSliceOp->getResult(0));
+    if (failed(dstMemref))
+      return failure();
 
     // Take a subview of the dst.
     auto dstMemrefType = dstMemref->getType().cast<MemRefType>();
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             insertSliceOp.getSourceType().getRank(), dstMemrefType,
             insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
             insertSliceOp.getMixedStrides())

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     assert(writeOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // Create a new transfer_write on buffer that doesn't have a return value.
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
     FailureOr<Value> resultBuffer =
         state.getResultBuffer(rewriter, op->getResult(0));
+    if (failed(resultBuffer))
+      return failure();
     rewriter.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.vector(), *resultBuffer, writeOp.indices(),
         writeOp.permutation_mapAttr(), writeOp.in_boundsAttr());
     replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
 
     return success();
   }
 };

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

 } // namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadTensorOpPattern>(moduleOp.getContext());
   (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
 }
 
-static Optional<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
+static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                 MemRefType type,
                                                 ArrayRef<Value> dynShape) {
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
   auto options = std::make_unique<BufferizationOptions>();
   if (useAlloca) {