[mlir][bufferize] Better propagation of errors

[springerm]

Summary

Return immediately when an op bufferization pattern fails.

[springerm]

update

[springerm]

update

[springerm]

update

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

let methods = [
           Regions of an op should be inlined into the new op instead of cloning
           them. This is not only more efficient, but also necessary so that no
           analysis results are lost. (Bufferization decisions are tracked via
           OpOperand pointers and cloned ops have new OpOperands.) If regions are
           cloned instead of inlined, additional buffer copies may be inserted.
 
           This method will never be called on ops that do not have at least one
           tensor operand/result.
+
+          The return value of this method indicates whether there was an error
+          while bufferizing this op (such as failing to create a new buffer
+          allocation op). The bufferization driver immediately stops bufferizing
+          the input IR and returns `failure` in that case. If this op is
+          expected to survive bufferization, `success` should be returned
+          (together with `allow-unknown-ops` enabled).
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "RewriterBase &":$rewriter,
                       "BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

LogicalResult bufferize(RewriterBase &rewriter,
                             BufferizationState &state) const {
       // to_tensor cannot be bufferized. However, other ops that are using
       // to_tensor's result will eventually be bufferized. At that point, they
       // will start using to_tensor's memref operand. Once all users of
       // to_tensor are bufferized, the op will not have any users anymore and
       // DCE away. In case of partial bufferization, to_memref(to_tensor(x))
       // constructs may be left over. These are folded by the canonicalizer or
       // FinalizingBufferize.
-      return failure();
+      return success();
     }
 
     bool isWritable(Value value, const AnalysisState &state) const {
       // It is unknown whether the memref operand is writable or not.
       return false;
     }
   }];
 

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

void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast, TensorLoadToMemref>(
       context);
 }
 
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
                                     BufferizationState &state) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
-  return foldToMemrefToTensorPair(rewriter, *this);
+  (void)foldToMemrefToTensorPair(rewriter, *this);
+  // Note: The return value of `bufferize` indicates whether there was an error
+  // or not. (And not whether the pattern matched or not.)
+  return success();
 }
 
 Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {
   return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
       .getOperation();
 }
 
 Optional<Value> CloneOp::buildClone(OpBuilder &builder, Value alloc) {

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

for (unsigned i = 0; i < worklist.size(); ++i) {
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       continue;
     // Skip ops that no longer have tensor semantics.
     if (!hasTensorSemantics(op))
       continue;
     // Bufferize the op.
     rewriter.setInsertionPoint(op);
-    (void)bufferizableOp.bufferize(rewriter, bufferizationState);
+    if (failed(bufferizableOp.bufferize(rewriter, bufferizationState)))
+      return op->emitError("failed to bufferize op");
   }
 
   // Fold all to_memref(to_tensor(x)) pairs.
   for (Operation *op : toMemrefOps) {
     rewriter.setInsertionPoint(op);
     (void)bufferization::foldToMemrefToTensorPair(rewriter,
                                                   cast<ToMemrefOp>(op));
   }

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
 #ifndef NDEBUG
     auto returnOp = cast<func::ReturnOp>(op);
     assert(isa<FuncOp>(returnOp->getParentOp()) &&
            "only support FuncOp parent for ReturnOp");
 #endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
-    return failure();
+    return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
   /// Rewrite function bbArgs and return values into buffer form. This function
   /// bufferizes the function signature and the ReturnOp. When the entire
   /// function body has been bufferized, function return types can be switched

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="bufferize-function-boundaries=1" -split-input-file -verify-diagnostics
 
 func.func private @foo() -> tensor<?xf32>
 
 func.func @bar() -> tensor<?xf32> {
   %foo = constant @foo : () -> (tensor<?xf32>)
 // expected-error @+1 {{expected a CallOp}}
   %res = call_indirect %foo() : () -> (tensor<?xf32>)
   return %res : tensor<?xf32>
 }
 
 // -----
 
-// expected-error @+2 {{op was not bufferized}}
-// expected-error @+1 {{cannot bufferize bodiless function that returns a tensor}}
+// expected-error @+2 {{cannot bufferize bodiless function that returns a tensor}}
+// expected-error @+1 {{failed to bufferize op}}
 func.func private @foo() -> tensor<?xf32>
 
 // -----
 
 // expected-error @+1 {{cannot bufferize a FuncOp with tensors and without a unique ReturnOp}}
 func.func @swappy(%cond1 : i1, %cond2 : i1, %t1 : tensor<f32>, %t2 : tensor<f32>)
     -> (tensor<f32>, tensor<f32>)
 {
func.func @to_memref_op_is_writing(
   %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
   %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
 
   return %r1, %r2 : vector<5xf32>, vector<5xf32>
 }
 
 // -----
 
-// expected-error @+2 {{op was not bufferized}}
+// expected-error @+2 {{failed to bufferize op}}
 // expected-error @+1 {{cannot bufferize bodiless function that returns a tensor}}
 func.func private @foo(%t : tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
 
 func.func @call_to_unknown_tensor_returning_func(%t : tensor<?xf32>) {
   call @foo(%t) : (tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
   return
 }