[mlir] Add a hook to PatternRewriter to allow for patterns to notify why a match failed.

[rriddle]

Summary

This revision adds a new hook, notifyMatchFailure, that allows for notifying the rewriter that a match failure is coming with the provided reason. This hook takes as a parameter a callback that fills a Diagnostic instance with the reason why the match failed. This allows for the rewriter to decide how this information can be displayed to the end-user, and may completely ignore it if desired(opt mode). For now, DialectConversion is updated to include this information in the debug output.

Depends On D76202

[jpienaar]

mlir/lib/Transforms/DialectConversion.cpp

996

Is a match failing really an error?

[rriddle]

Resolve comments

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

996

Changed to Remark. Given that we don't include the severity in the message here, it doesn't matter what it gets set to.

[jpienaar]

Looks good, then we can generate mismatch messages from DRR and resolve a long open FR :-)

[mehdi_amini]

mlir/lib/Transforms/DialectConversion.cpp

1001

If this was in a header, we could try to rely on inlining to make it purely zero cost in release mode at call-sites by eliminating all the debug strings. Right now they would still be in the binary I think (LTO may help). Have you looked into this possibility?

mlir/include/mlir/IR/PatternMatch.h

public:
   /// callable.
   template <typename CallableT>
   void updateRootInPlace(Operation *root, CallableT &&callable) {
     startRootUpdate(root);
     callable();
     finalizeRootUpdate(root);
   }
 
+  /// Notify the pattern rewriter that the pattern is failing to match the given
+  /// operation, and provide a callback to populate a diagnostic with the reason
+  /// why the failure occurred. This method allows for derived rewriters to
+  /// optionally hook into the reason why a pattern failed, and display it to
+  /// users.
+  virtual LogicalResult
+  notifyMatchFailure(Operation *op,
+                     function_ref<void(Diagnostic &)> reasonCallback) {
+    return failure();
+  }
+  LogicalResult notifyMatchFailure(Operation *op, const Twine &msg) {
+    return notifyMatchFailure(op, [&](Diagnostic &diag) { diag << msg; });
+  }
+  LogicalResult notifyMatchFailure(Operation *op, const char *msg) {
+    return notifyMatchFailure(op, Twine(msg));
+  }
+
 protected:
   explicit PatternRewriter(MLIRContext *ctx) : OpBuilder(ctx) {}
   virtual ~PatternRewriter();
 
   // These are the callback methods that subclasses can choose to implement if
   // they would like to be notified about certain types of mutations.
 
   /// Notify the pattern rewriter that the specified operation is about to be

mlir/include/mlir/Transforms/DialectConversion.h

public:
   void startRootUpdate(Operation *op) override;
 
   /// PatternRewriter hook for updating the root operation in-place.
   void finalizeRootUpdate(Operation *op) override;
 
   /// PatternRewriter hook for updating the root operation in-place.
   void cancelRootUpdate(Operation *op) override;
 
+  /// PatternRewriter hook for notifying match failure reasons.
+  LogicalResult
+  notifyMatchFailure(Operation *op,
+                     function_ref<void(Diagnostic &)> reasonCallback) override;
+  using PatternRewriter::notifyMatchFailure;
+
   /// Return a reference to the internal implementation.
   detail::ConversionPatternRewriterImpl &getImpl();
 
 private:
   std::unique_ptr<detail::ConversionPatternRewriterImpl> impl;
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Transforms/DialectConversion.cpp

 #endif
   // Erase the last update for this operation.
   auto stateHasOp = [op](const auto &it) { return it.getOperation() == op; };
   auto &rootUpdates = impl->rootUpdates;
   auto it = llvm::find_if(llvm::reverse(rootUpdates), stateHasOp);
   rootUpdates.erase(rootUpdates.begin() + (rootUpdates.rend() - it));
 }
 
+/// PatternRewriter hook for notifying match failure reasons.
+LogicalResult ConversionPatternRewriter::notifyMatchFailure(
+    Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
+  LLVM_DEBUG({
+    Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark);
+    reasonCallback(diag);
+    impl->logger.startLine() << "** Failure : " << diag.str() << "\n";
+  });
+  return failure();
+}
+
 /// Return a reference to the internal implementation.
 detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() {
   return *impl;
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion Patterns
 //===----------------------------------------------------------------------===//

mlir/test/lib/TestDialect/TestPatterns.cpp

 struct TestChangeProducerTypeF32ToF64 : public ConversionPattern {
   TestChangeProducerTypeF32ToF64(MLIRContext *ctx)
       : ConversionPattern("test.type_producer", 1, ctx) {}
   PatternMatchResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     // If the type is F32, change the type to F64.
     if (!Type(*op->result_type_begin()).isF32())
-      return matchFailure();
+      return rewriter.notifyMatchFailure(op, "expected single f32 operand");
     rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, rewriter.getF64Type());
     return matchSuccess();
   }
 };
 struct TestChangeProducerTypeF32ToInvalid : public ConversionPattern {
   TestChangeProducerTypeF32ToInvalid(MLIRContext *ctx)
       : ConversionPattern("test.type_producer", 10, ctx) {}
   PatternMatchResult