[mlir] Fail early if AnalysisState::getBuffer() returns failure

[ashay-github]

Summary

This patch updates calls to AnalysisState::getBuffer() so that we return
early with a failure if the call does not succeed.

[springerm]

nice cleanup thx!

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     RankedTensorType tensorResultType = collapseShapeOp.getResultType();
     OpOperand &srcOperand = collapseShapeOp->getOpOperand(0) /*src*/;
     auto bufferType = state.getBufferType(srcOperand).cast<MemRefType>();
 
     if (tensorResultType.getRank() == 0) {
       // 0-d collapses must go through a different op builder.
-      Value buffer = *state.getBuffer(rewriter, srcOperand);
+      auto buffer = state.getBuffer(rewriter, srcOperand);
+      if (failed(buffer))
+        return failure();
       MemRefType resultType;
 
       if (bufferType.getLayout().isIdentity()) {
         // Standard layout: result type has no offset.
         MemRefLayoutAttrInterface layout;
         resultType = MemRefType::get({}, tensorResultType.getElementType(),
                                      layout, bufferType.getMemorySpace());
       } else {
         // Source memref has a layout map: result type has the same offset as
         // the source type.
         SmallVector<int64_t> strides;
         int64_t offset;
         if (failed(getStridesAndOffset(bufferType, strides, offset)))
           return failure();
         AffineMap resultLayout =
             makeStridedLinearLayoutMap({}, offset, op->getContext());
         resultType =
             MemRefType::get({}, tensorResultType.getElementType(), resultLayout,
                             bufferType.getMemorySpaceAsInt());
       }
 
       replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
-          rewriter, op, resultType, buffer, collapseShapeOp.reassociation());
+          rewriter, op, resultType, *buffer, collapseShapeOp.reassociation());
       return success();
     }
 
     // If the dims are not collapsible (due to an incompatible source layout
     // map), force an out-of-place bufferization, i.e., a buffer copy. This
     // newly allocated buffer will have no layout map and thus be collapsible.
     bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
         bufferType, collapseShapeOp.getReassociationIndices());
     Optional<BufferizationState::ForceInPlacability> overrideInPlace =
         canBeCollapsed
             ? None
             : Optional<BufferizationState::ForceInPlacability>(
                   BufferizationState::ForceInPlacability::FORCE_OUT_OF_PLACE);
-    Value buffer = *state.getBuffer(rewriter, srcOperand, overrideInPlace);
+    auto buffer = state.getBuffer(rewriter, srcOperand, overrideInPlace);
+    if (failed(buffer))
+      return failure();
 
     // Result type is inferred by the builder.
     replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
-        rewriter, op, buffer, collapseShapeOp.getReassociationIndices());
+        rewriter, op, *buffer, collapseShapeOp.getReassociationIndices());
     return success();
   }
 };
 
 /// Bufferization of tensor.dim. Replace with memref.dim.
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
struct DimOpInterface
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
-    Value v = *state.getBuffer(rewriter, dimOp->getOpOperand(0) /*source*/);
-    replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, v, dimOp.index());
+    auto v = state.getBuffer(rewriter, dimOp->getOpOperand(0) /*source*/);
+    if (failed(v))
+      return failure();
+    replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, *v,
+                                                dimOp.index());
     return success();
   }
 };
 
 /// Bufferization of tensor.expand_shape. Replace with memref.expand_shape.
 struct ExpandShapeOpInterface
     : public BufferizableOpInterface::ExternalModel<ExpandShapeOpInterface,
                                                     tensor::ExpandShapeOp> {
BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     auto tensorResultType = expandShapeOp.getResultType();
-    Value buffer =
-        *state.getBuffer(rewriter, expandShapeOp->getOpOperand(0) /*src*/);
+    auto buffer =
+        state.getBuffer(rewriter, expandShapeOp->getOpOperand(0) /*src*/);
+    if (failed(buffer))
+      return failure();
 
     // Memref result type is inferred by the builder based on reassociation
     // indices and result shape.
     replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
-        rewriter, op, tensorResultType.getShape(), buffer,
+        rewriter, op, tensorResultType.getShape(), *buffer,
         expandShapeOp.getReassociationIndices());
     return success();
   }
 };
 
 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
 struct ExtractSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
struct ExtractSliceOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
 
     // Even if this op was decided to bufferize out-of-place, do not insert the
     // buffer copy yet. This is done later in this function.
-    Value srcMemref =
-        *state.getBuffer(rewriter, extractSliceOp->getOpOperand(0) /*source*/,
-                         BufferizationState::ForceInPlacability::FORCE_INPLACE);
-    auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
+    auto srcMemref =
+        state.getBuffer(rewriter, extractSliceOp->getOpOperand(0) /*source*/,
+    if (failed(srcMemref))
+      return failure();
+    auto srcMemrefType = srcMemref->getType().cast<MemRefType>();
     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 =
           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();
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     OffsetSizeAndStrideOpInterface::expandToRank(
-        srcMemref, mixedOffsets, mixedSizes, mixedStrides,
+        *srcMemref, mixedOffsets, mixedSizes, mixedStrides,
         [&](Value target, int64_t dim) -> OpFoldResult {
           auto shapedType = target.getType().cast<ShapedType>();
           if (shapedType.isDynamicDim(dim))
             return rewriter.create<memref::DimOp>(loc, target, dim).result();
           return rewriter.getIndexAttr(shapedType.getDimSize(dim));
         });
     // Bufferize to subview.
     auto subviewMemRefType = memref::SubViewOp::inferRankReducedResultType(
                                  dstTensorType.getRank(), srcMemrefType,
                                  mixedOffsets, mixedSizes, mixedStrides)
                                  .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
-        loc, subviewMemRefType, srcMemref, mixedOffsets, mixedSizes,
+        loc, subviewMemRefType, *srcMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // If not inplaceable, copy.
     if (!inplace) {
       // Do not copy if the copied data is never read.
       if (state.getAnalysisState().isValueRead(extractSliceOp.result()))
         if (failed(createMemCpy(rewriter, extractSliceOp.getLoc(), subView,
                                 alloc, state.getOptions())))
struct ExtractOpInterface
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
-    Value srcMemref =
-        *state.getBuffer(rewriter, extractOp->getOpOperand(0) /*tensor*/);
-    replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, srcMemref,
+    auto srcMemref =
+        state.getBuffer(rewriter, extractOp->getOpOperand(0) /*tensor*/);
+    if (failed(srcMemref))
+      return failure();
+    replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, *srcMemref,
                                                  extractOp.indices());
     return success();
   }
 };
 
 // Implements backtracking to traverse indices of the output buffer while
 // iterating over op.elements().
 static void createStores(RewriterBase &rewriter, Location loc, int dim,
auto subviewMemRefType =
             mixedOffsets, mixedSizes, mixedStrides)
             .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // Copy tensor. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
-    Value srcMemref =
-        *state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
-    if (failed(createMemCpy(rewriter, loc, srcMemref, subView,
-                            state.getOptions())))
+    auto srcMemref =
+        state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
+    if (failed(srcMemref) || failed(createMemCpy(rewriter, loc, *srcMemref,
+                                                 subView, state.getOptions())))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
 /// Bufferization of tensor.rank. Replace with memref.rank.
struct RankOpInterface
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto rankOp = cast<tensor::RankOp>(op);
-    Value v = *state.getBuffer(rewriter, rankOp->getOpOperand(0) /*source*/);
+    auto v = state.getBuffer(rewriter, rankOp->getOpOperand(0) /*source*/);
+    if (failed(v))
+      return failure();
     replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
-                                                 v);
+                                                 *v);
     return success();
   }
 };
 
 /// Bufferization of tensor.reshape. Replace with memref.reshape.
 struct ReshapeOpInterface
     : public BufferizableOpInterface::ExternalModel<ReshapeOpInterface,
                                                     tensor::ReshapeOp> {