[mlir] Support collecting logs from notifyMatchFailure().

[Chia-hungDuan]

Summary

Let the user registers their own handler to processing the matching
failure information.

[jpienaar]

Why isn't all the methods guarded behind LLVM_DEBUG (or similar debug macro/define)? Seems like that would remove chance of unexpected behavior and make it explicitly for debugging. Downside is then folks can't enable debugging without using the same guarding - but I think that makes expectations clearer.

[mehdi_amini]

In D110896#3036503, @jpienaar wrote:

Why isn't all the methods guarded behind LLVM_DEBUG (or similar debug macro/define)? Seems like that would remove chance of unexpected behavior and make it explicitly for debugging. Downside is then folks can't enable debugging without using the same guarding - but I think that makes expectations clearer.

In headers, it would have to be #if LLVM_ENABLE_ABI_BREAKING_CHECKS or something like that.

[Chia-hungDuan]

In D110896#3036503, @jpienaar wrote:

Why isn't all the methods guarded behind LLVM_DEBUG (or similar debug macro/define)? Seems like that would remove chance of unexpected behavior and make it explicitly for debugging. Downside is then folks can't enable debugging without using the same guarding - but I think that makes expectations clearer.

In D110896#3036917, @mehdi_amini wrote:

In D110896#3036503, @jpienaar wrote:

Why isn't all the methods guarded behind LLVM_DEBUG (or similar debug macro/define)? Seems like that would remove chance of unexpected behavior and make it explicitly for debugging. Downside is then folks can't enable debugging without using the same guarding - but I think that makes expectations clearer.

In headers, it would have to be #if LLVM_ENABLE_ABI_BREAKING_CHECKS or something like that.

Let me try if I can turn these all into debug build only gracefully.

[Chia-hungDuan]

Code simplify

[rriddle]

Can you provide more detail on what problem you are trying to solve here? and how this diagnostics support is supposed to compose/be used in then general infra? The current description and the commit itself leads me to believe this is trying to build a general solution for a dialect conversion specific problem. If the problem is just for dialect conversion, there are likely more local solutions that would work as well.

[Chia-hungDuan]

In D110896#3041236, @rriddle wrote:

Can you provide more detail on what problem you are trying to solve here? and how this diagnostics support is supposed to compose/be used in then general infra? The current description and the commit itself leads me to believe this is trying to build a general solution for a dialect conversion specific problem. If the problem is just for dialect conversion, there are likely more local solutions that would work as well.

The user wants to create a report for a failed dialect conversion. They would like to collect the failure reason generated during dialect conversion (match failure is the main one) then they can have a report which tells the reasons in detail that why an operation can't be converted. In the usage, they will have a debug build tool to run over those failed cases again so that they will have debugging message there. And they would like to have a solution that they can capture the message while conversion so that they don't need to parse/classify the logs.

I think this approach is lightweight and yes, it's a general solution which I think it avoids adding some hooks like passing callback for certain path. What do you think?

[Chia-hungDuan]

In D110896#3041262, @Chia-hungDuan wrote:

In D110896#3041236, @rriddle wrote:

Can you provide more detail on what problem you are trying to solve here? and how this diagnostics support is supposed to compose/be used in then general infra? The current description and the commit itself leads me to believe this is trying to build a general solution for a dialect conversion specific problem. If the problem is just for dialect conversion, there are likely more local solutions that would work as well.

The user wants to create a report for a failed dialect conversion. They would like to collect the failure reason generated during dialect conversion (match failure is the main one) then they can have a report which tells the reasons in detail that why an operation can't be converted. In the usage, they will have a debug build tool to run over those failed cases again so that they will have debugging message there. And they would like to have a solution that they can capture the message while conversion so that they don't need to parse/classify the logs.

I think this approach is lightweight and yes, it's a general solution which I think it avoids adding some hooks like passing callback for certain path. What do you think?

One thing to add, user can achieve this by using ScopedDiagnosticHandler as well but I'm thinking if we can provide one like this then they don't need to worry about how to keep the original diagnostic handler work as usual.

[Chia-hungDuan]

Add a way to group the diagnostics. It means we can group the related
diagnostics together to produce better debugging message.

[rriddle]

How do you see all of this being used outside of the dialect conversion use case? The point I made earlier was whether we should start with something that is passed directly to dialect conversion instead of changing the general diagnostics API. If this is something that is only serving a specific use case, we should keep the scope of it limited.

[Chia-hungDuan]

In D110896#3105051, @rriddle wrote:

How do you see all of this being used outside of the dialect conversion use case? The point I made earlier was whether we should start with something that is passed directly to dialect conversion instead of changing the general diagnostics API. If this is something that is only serving a specific use case, we should keep the scope of it limited.

I thought last time I mentioned that passing the callback in dialect conversion path it seems to be hacky and you mentioned that accessing the diagnostic may be useful(sorry, correct me if I misunderstand). Then we talked if we want to let the user use this, it'll be better to make this more useful (last change only gave a handler wrapper), like if we can let them know the set of diagnostics which are associated with a match failure of certain operation.

So far, the case is only focusing on collecting the message output during pattern matching. I'm not sure if they would like to collect more information other than dialect conversion. Thus I tried to make this for general path and keep as little as impact in normal path.

If this is not a right way, I will suggest the user to report the diagnostics(not through rewriter.notifyMatchFailure()) in their patterns and register the handler to read them. Except the patterns written in DRR, others can be handled well.

[rriddle]

In D110896#3105057, @Chia-hungDuan wrote:

In D110896#3105051, @rriddle wrote:

How do you see all of this being used outside of the dialect conversion use case? The point I made earlier was whether we should start with something that is passed directly to dialect conversion instead of changing the general diagnostics API. If this is something that is only serving a specific use case, we should keep the scope of it limited.

I thought last time I mentioned that passing the callback in dialect conversion path it seems to be hacky and you mentioned that accessing the diagnostic may be useful(sorry, correct me if I misunderstand). Then we talked if we want to let the user use this, it'll be better to make this more useful (last change only gave a handler wrapper), like if we can let them know the set of diagnostics which are associated with a match failure of certain operation.

So far, the case is only focusing on collecting the message output during pattern matching. I'm not sure if they would like to collect more information other than dialect conversion. Thus I tried to make this for general path and keep as little as impact in normal path.

If this is not a right way, I will suggest the user to report the diagnostics(not through rewriter.notifyMatchFailure()) in their patterns and register the handler to read them. Except the patterns written in DRR, others can be handled well.

I see, seems like some wires got crossed. Let's sync this week to discuss it.

[Chia-hungDuan]

Per our discussion, add a specific way to support collecting the matching failure message for now
Only add it in applyAnalysisConversion because the user will only use it on that.

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

501–502

Cab we initialize it in the impl separately? Having it here leaks the implementation details a bit. (This constructor really shouldn't be public anyways).

911

Can you document the callback here? What invariants is it supposed to provide? How would a potential user interact with it/know when they would want to provide it? Also, do we only want the diagnostic here? What about the pattern/op being matched/etc.? What information is useful to provide back?

mlir/lib/Transforms/Utils/DialectConversion.cpp

1487

This isn't the same as LLVM_DEBUG. LLVM_DEBUG also handles checking if the current debug mode enabled is dialect conversion. Can we just have a bool that is defaulted to true/false depending on notifyCallback, that gets explicitly set to true inside of an LLVM_DEBUG code block?

[Chia-hungDuan]

Address review comment

mlir/include/mlir/Transforms/DialectConversion.h

911

Also added the log of the pattern and operation. According to the current architecture, The match failure reason will come before pattern name, it'll be like,

match failure ....
Failed to apply pattern 'Foo' on BarOp

mlir/lib/Transforms/Utils/DialectConversion.cpp

1487

Reverted back to use LLVM_DEBUG I think that's enough for the initial support

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

907–910

Is this only intended for analysis conversions, or should we do this for partial/full conversions too?

908–910

mlir/lib/Transforms/Utils/DialectConversion.cpp

1497–1508

nit: I would just inline all of these into the class.

1983–1988

What does the pattern debug log output look like with this?

[Chia-hungDuan]

The messages captured will be looks like

// The failure reason of applying a certain pattern, pattern name comes after this
expected single f32 operand                                                         
                                                                                    
// The pattern name and the operation it tried to apply
Failed to apply pattern "(anonymous namespace)::TestChangeProducerTypeF32ToF64" on op: 
%0 = "test.type_producer"() : () -> bf16

mlir/include/mlir/Transforms/DialectConversion.h

907–910

Per discussion, now we will only have it in analysis conversions.

mlir/lib/Transforms/Utils/DialectConversion.cpp

1983–1988

As you mentioned, the curly brackets may have little problem. Added a new line for output. Now it looks like

Failed to apply pattern "(anonymous namespace)::TestRegionRewriteBlockMovement" on op:
"test.region"() ( {
^bb0(%arg0: i64):  // no predecessors
  br ^bb1(%arg0 : i64)
^bb1(%0: i64):  // pred: ^bb0
  "test.invalid"(%0) : (i64) -> ()
}) {legalizer.erase_old_blocks, legalizer.should_clone} : () -> ()

[Chia-hungDuan]

Address review comment

mlir/include/mlir/Transforms/DialectConversion.h

 struct ConversionPatternRewriterImpl;
 } // end namespace detail
 
 /// This class implements a pattern rewriter for use with ConversionPatterns. It
 /// extends the base PatternRewriter and provides special conversion specific
 /// hooks.
 class ConversionPatternRewriter final : public PatternRewriter {
 public:
-  ConversionPatternRewriter(MLIRContext *ctx);
+  explicit ConversionPatternRewriter(MLIRContext *ctx);
   ~ConversionPatternRewriter() override;
 
   /// Apply a signature conversion to the entry block of the given region. This
   /// replaces the entry block with a new block containing the updated
   /// signature. The new entry block to the region is returned for convenience.
   ///
   /// If provided, `converter` will be used for any materializations.
   Block *
   applySignatureConversion(Region *region,
                            TypeConverter::SignatureConversion &conversion,

 
 /// Apply an analysis conversion on the given operations, and all nested
 /// operations. This method analyzes which operations would be successfully
 /// converted to the target if a conversion was applied. All operations that
 /// were found to be legalizable to the given 'target' are placed within the
 /// provided 'convertedOps' set; note that no actual rewrites are applied to the
 /// operations on success and only pre-existing operations are added to the set.
 /// This method only returns failure if there are unreachable blocks in any of
-/// the regions nested within 'ops'.
-LogicalResult applyAnalysisConversion(ArrayRef<Operation *> ops,
-                                      ConversionTarget &target,
-                                      const FrozenRewritePatternSet &patterns,
-                                      DenseSet<Operation *> &convertedOps);
-LogicalResult applyAnalysisConversion(Operation *op, ConversionTarget &target,
-                                      const FrozenRewritePatternSet &patterns,
-                                      DenseSet<Operation *> &convertedOps);
+/// the regions nested within 'ops'. There's an additional argument
+/// `notifyCallback` which is used for collecting the conversion diagnostic
+/// generated during the conversion. Most diagnostics reported to this callback
+/// may only be available in debug mode.
+LogicalResult applyAnalysisConversion(
+    ArrayRef<Operation *> ops, ConversionTarget &target,
+    DenseSet<Operation *> &convertedOps,
+    function_ref<void(Diagnostic &)> notifyCallback = nullptr);
+LogicalResult applyAnalysisConversion(
+    Operation *op, ConversionTarget &target,
+    DenseSet<Operation *> &convertedOps,
+    function_ref<void(Diagnostic &)> notifyCallback = nullptr);
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_DIALECTCONVERSION_H_

mlir/lib/Transforms/Utils/DialectConversion.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriterImpl
 //===----------------------------------------------------------------------===//
 namespace mlir {
 namespace detail {
 struct ConversionPatternRewriterImpl {
-  ConversionPatternRewriterImpl(PatternRewriter &rewriter)
-      : argConverter(rewriter, unresolvedMaterializations) {}
+  explicit ConversionPatternRewriterImpl(PatternRewriter &rewriter)
+      : argConverter(rewriter, unresolvedMaterializations),
+        notifyCallback(nullptr) {}
 
   /// Cleanup and destroy any generated rewrite operations. This method is
   /// invoked when the conversion process fails.
   void discardRewrites();
 
   /// Apply all requested operation rewrites. This method is invoked when the
   /// conversion process succeeds.
   void applyRewrites();
struct ConversionPatternRewriterImpl {
   void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
                                    Location origRegionLoc);
 
   /// Notifies that a pattern match failed for the given reason.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback);
 
+  /// Registers a callback to get notified when there's any pattern match
+  /// failure.
+  void setNotifyCallback(function_ref<void(Diagnostic &)> callback);
+
+  /// Check if there's a registered callback.
+  bool hasNotifyCallback() const;
+
+  /// Invoke the callback with the given diagnostic.
+  void notifyMatchResult(Diagnostic &diag);
+
   //===--------------------------------------------------------------------===//
   // State
   //===--------------------------------------------------------------------===//
 
   // Mapping between replaced values that differ in type. This happens when
   // replacing a value with one of a different type.
   ConversionValueMapping mapping;
 
struct ConversionPatternRewriterImpl {
   /// with values with different result types than the original operation, e.g.
   /// 1->N conversion of some kind.
   SmallVector<unsigned, 4> operationsWithChangedResults;
 
   /// The current type converter, or nullptr if no type converter is currently
   /// active.
   TypeConverter *currentTypeConverter = nullptr;
 
+  /// This allows the user to collect the match failure message.
+  function_ref<void(Diagnostic &)> notifyCallback;
+
 #ifndef NDEBUG
   /// A set of operations that have pending updates. This tracking isn't
   /// strictly necessary, and is thus only active during debug builds for extra
   /// verification.
   SmallPtrSet<Operation *, 1> pendingRootUpdates;
 
   /// A logger used to emit diagnostics during the conversion process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
   // This original region has already had its conversion set computed, so there
   // shouldn't be any new failures.
   (void)result;
   assert(succeeded(result) && "expected region to have no unreachable blocks");
 }
 
 LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
+    if (hasNotifyCallback())
+      notifyMatchResult(diag);
   });
   return failure();
 }
 
+void ConversionPatternRewriterImpl::setNotifyCallback(
+    function_ref<void(Diagnostic &)> callback) {
+  notifyCallback = callback;
+}
+
+bool ConversionPatternRewriterImpl::hasNotifyCallback() const {
+  return notifyCallback ? true : false;
+}
+
+void ConversionPatternRewriterImpl::notifyMatchResult(Diagnostic &diag) {
+  notifyCallback(diag);
+}
+
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriter
 //===----------------------------------------------------------------------===//
 
 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
     : PatternRewriter(ctx),
       impl(new detail::ConversionPatternRewriterImpl(*this)) {}
+
 ConversionPatternRewriter::~ConversionPatternRewriter() {}
 
 void ConversionPatternRewriter::replaceOpWithIf(
     Operation *op, ValueRange newValues, bool *allUsesReplaced,
     llvm::unique_function<bool(OpOperand &) const> functor) {
   // TODO: To support this we will need to rework a bit of how replacements are
   // tracked, given that this isn't guranteed to replace all of the uses of an
   // operation. The main change is that now an operation can be replaced
OperationLegalizer::legalizeWithPattern(Operation *op,
   // Functor that returns if the given pattern may be applied.
   auto canApply = [&](const Pattern &pattern) {
     return canApplyPattern(op, pattern, rewriter);
   };
 
   // Functor that cleans up the rewriter state after a pattern failed to match.
   RewriterState curState = rewriterImpl.getCurrentState();
   auto onFailure = [&](const Pattern &pattern) {
-    LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern failed to match"));
+    LLVM_DEBUG({
+      logFailure(rewriterImpl.logger, "pattern failed to match");
+      if (rewriterImpl.hasNotifyCallback()) {
+        Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark);
+        diag << "Failed to apply pattern \"" << pattern.getDebugName()
+             << "\" on " << *op;
+        rewriterImpl.notifyMatchResult(diag);
+      }
+    });
     rewriterImpl.resetState(curState);
     appliedPatterns.erase(&pattern);
   };
 
   // Functor that performs additional legalization when a pattern is
   // successfully applied.
   auto onSuccess = [&](const Pattern &pattern) {
     auto result = legalizePatternResult(op, pattern, rewriter, curState);

 struct OperationConverter {
   explicit OperationConverter(ConversionTarget &target,
                               const FrozenRewritePatternSet &patterns,
                               OpConversionMode mode,
                               DenseSet<Operation *> *trackedOps = nullptr)
       : opLegalizer(target, patterns), mode(mode), trackedOps(trackedOps) {}
 
   /// Converts the given operations to the conversion target.
-  LogicalResult convertOperations(ArrayRef<Operation *> ops);
+  LogicalResult
+  convertOperations(ArrayRef<Operation *> ops,
+                    function_ref<void(Diagnostic &)> notifyCallback = nullptr);
 
 private:
   /// Converts an operation with the given rewriter.
   LogicalResult convert(ConversionPatternRewriter &rewriter, Operation *op);
 
   /// This method is called after the conversion process to legalize any
   /// remaining artifacts and complete the conversion.
   LogicalResult finalize(ConversionPatternRewriter &rewriter);
if (failed(opLegalizer.legalize(op, rewriter))) {
     // Analysis conversions don't fail if any operations fail to legalize,
     // they are only interested in the operations that were successfully
     // legalized.
     trackedOps->insert(op);
   }
   return success();
 }
 
-LogicalResult OperationConverter::convertOperations(ArrayRef<Operation *> ops) {
+LogicalResult OperationConverter::convertOperations(
+    ArrayRef<Operation *> ops,
+    function_ref<void(Diagnostic &)> notifyCallback) {
   if (ops.empty())
     return success();
   ConversionTarget &target = opLegalizer.getTarget();
 
   // Compute the set of operations and blocks to convert.
   SmallVector<Operation *> toConvert;
   for (auto *op : ops) {
     toConvert.emplace_back(op);
     for (auto &region : op->getRegions())
       if (failed(computeConversionSet(region.getBlocks(), region.getLoc(),
                                       toConvert, &target)))
         return failure();
   }
 
   // Convert each operation and discard rewrites on failure.
   ConversionPatternRewriter rewriter(ops.front()->getContext());
   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
+  rewriterImpl.setNotifyCallback(notifyCallback);
+
   for (auto *op : toConvert)
     if (failed(convert(rewriter, op)))
       return rewriterImpl.discardRewrites(), failure();
 
   // Now that all of the operations have been converted, finalize the conversion
   // process to ensure any lingering conversion artifacts are cleaned up and
   // legalized.
   if (failed(finalize(rewriter)))

 
 //===----------------------------------------------------------------------===//
 // Analysis Conversion
 
 LogicalResult
 mlir::applyAnalysisConversion(ArrayRef<Operation *> ops,
                               ConversionTarget &target,
                               const FrozenRewritePatternSet &patterns,
-                              DenseSet<Operation *> &convertedOps) {
+                              DenseSet<Operation *> &convertedOps,
+                              function_ref<void(Diagnostic &)> notifyCallback) {
   OperationConverter opConverter(target, patterns, OpConversionMode::Analysis,
                                  &convertedOps);
-  return opConverter.convertOperations(ops);
+  return opConverter.convertOperations(ops, notifyCallback);
 }
 LogicalResult
 mlir::applyAnalysisConversion(Operation *op, ConversionTarget &target,
                               const FrozenRewritePatternSet &patterns,
-                              DenseSet<Operation *> &convertedOps) {
+                              DenseSet<Operation *> &convertedOps,
+                              function_ref<void(Diagnostic &)> notifyCallback) {
   return applyAnalysisConversion(llvm::makeArrayRef(op), target, patterns,
-                                 convertedOps);
+                                 convertedOps, notifyCallback);
 }