[Mlir-commits] [mlir] 0b293bf - [mlir][bufferize] Better propagation of errors

[Matthias Springer]

Author: Matthias Springer
Date: 2022-05-16T23:17:01+02:00
New Revision: 0b293bf0451cd695239867d8dac9b239ab601a70

URL: https://github.com/llvm/llvm-project/commit/0b293bf0451cd695239867d8dac9b239ab601a70
DIFF: https://github.com/llvm/llvm-project/commit/0b293bf0451cd695239867d8dac9b239ab601a70.diff

LOG: [mlir][bufferize] Better propagation of errors

Return immediately when an op bufferization patterns fails.

Differential Revision: https://reviews.llvm.org/D125087

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
    mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
    mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
    mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
    mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
index 620fff5506495..76b0659ae96eb 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -216,6 +216,13 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
 
           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",

diff  --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 5f1a132c7d48c..804a8bf4f7f57 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -133,7 +133,7 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
       // 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 {

diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 0e8b2614d47bb..6509a8ebd22d9 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -354,7 +354,10 @@ void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
 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) {

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 0727b1ce9da83..3bde5a59550e0 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -416,7 +416,8 @@ bufferization::bufferizeOp(Operation *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.

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index a6a6935bc18da..16989dff36722 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -389,7 +389,7 @@ struct ReturnOpInterface
 #endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
-    return failure();
+    return success();
   }
 };
 

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
index 8ab28773d7f7c..88a128921c6ec 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
@@ -11,8 +11,8 @@ func.func @bar() -> 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>
 
 // -----
@@ -262,7 +262,7 @@ func.func @to_memref_op_is_writing(
 
 // -----
 
-// 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)