[mlir] Add a new debug action framework.

[rriddle]

Summary

This revision adds the infrastructure for Debug Actions. This is a DEBUG only
API that allows for external entities to control various aspects of compiler
execution. This is conceptually similar to something like DebugCounters in LLVM, but at a lower level. This framework doesn't make any assumptions about how the higher level driver is controlling the execution, it merely provides a framework for connecting the two together. This means that on top of DebugCounter functionality, we could also provide more interesting drivers such as interactive execution. A high level overview of the workflow surrounding debug actions is
shown below:

• Compiler developer defines an action that is taken by the a pass, transformation, utility that they are developing.
• Depending on the needs, the developer dispatches various queries, pertaining to this action, to an action manager that will provide an answer as to what behavior the action should do.
• An external entity registers an action handler with the action manager, and provides the logic to resolve queries on actions.

The exact definition of an external entity is left opaque, to allow for more
interesting handlers.

This framework was proposed here: https://llvm.discourse.group/t/rfc-debug-actions-in-mlir-debug-counters-for-the-modern-world

[mehdi_amini]

Can you clang-format and rebase?

[rriddle]

Rebase

[rriddle]

In D84986#2426900, @mehdi_amini wrote:

Can you clang-format and rebase?

Done.

[rriddle]

Rebase

[jpienaar]

Overall looks good, thanks. I need to look at the follow up. I'd be very interested to see how we could hook this up with some user input/interactive stepping later :)

mlir/include/mlir/Support/DebugAction.h

132

Is the preferred method and the above the generic one? It would seem like the above could almost be private.

153

This is weird in the context of a failed status returned. Is FailureOr appropriate? E.g., is it really showing "handler registered and it returned true", so that if no handler registered then it runs. Perhaps a different type could help.

I mean, it failed to dispatch to handler as there is no handler, but also it could have been considered a success as there is no handler.

160

Is there a way to ensure all is compiled away outside debug mode? Maybe by having NDEBUG over larger sections/some code duplicated. That would make the NOP nature easier to see (at least for me :))

mlir/unittests/Support/DebugActionTest.cpp

62

This works quite nicely. Could you show one where either it runs up until 5th iteration? And a test for printing out the formatted arg strings?

[mehdi_amini]

mlir/docs/DebugActions.md

15

Something that wasn't clear to me when reading this was how would the developer define the queries. It turns out there is a pre-defined fixed list of queries that are available globally on the action manager for all possible actions.

[rriddle]

update

[rriddle]

mlir/docs/DebugActions.md

15

I want to relax the static aspect of this, but haven't yet finalized thoughts on how to generically+opaquely add external queries and keep everything nice and type safe. The section on the action manager mentions the set of available queries. Added more links here.

mlir/include/mlir/Support/DebugAction.h

132

This one will probably be used more frequently, but didn't want to prevent the above one if there is additional initialization necessary when creating the handler. For example, this is used in the debug counter in the child revision.

153

Discussed offline, keeping as FailureOr for now as that matches the diagnostic engine behavior. In a followup it may be interesting to add an UnhandledOr alias for FailureOr that better describes the real behavior.

160

Commented out the entire private section, and the internals of the public methods.

[mehdi_amini]

mlir/include/mlir/Support/DebugAction.h

216

I'm not sure about the right behavior for dispatchToHandler here.
Seems like the last registered handler will process the query and the others won't get it, does it generalize to every kind of query?
Wouldn't we want something that could log and let other handler manages the query as well?

I guess for a query that returns a value you need to figure what value to return, so the first handler that gets you a value may be reasonable. For void queries or any kind of tracing, it may not be the same.

255

What is the purpose of accepting this? Do you have examples of how it works?

In general the whole "arguments as string" is a bit weird to me.

[rriddle]

Update

[rriddle]

mlir/include/mlir/Support/DebugAction.h

216

I'm not sure yet either TBH. Each query would likely have it's own kind of reduction behavior, but picking the first one seems fine for me until we have a use case that really needs/wants something smarter.

255

Just removed opaque parameters for now.

[mehdi_amini]

mlir/include/mlir/Support/DebugAction.h

216

Right, maybe you can make mention of this in the doc?

[rriddle]

update

mlir/docs/DebugActions.md

+# Debug Actions
+
+[TOC]
+
+This file documents the infrastructure for `Debug Actions`. This is a DEBUG only
+API that allows for external entities to control various aspects of compiler
+execution. This is conceptually similar to something like `DebugCounters` in
+LLVM, but at a lower level. This framework doesn't make any assumptions about
+how the higher level driver is controlling the execution, it merely provides a
+framework for connecting the two together. A high level overview of the workflow
+surrounding debug actions is shown below:
+
+*   Compiler developer defines an `action` that is taken by the a pass,
+    transformation, utility that they are developing.
+*   Depending on the needs, the developer dispatches various queries, pertaining
+    to this action, to an `action manager` that will provide an answer as to
+    what behavior the action should do.
+*   An external entity registers an `action handler` with the action manager,
+    and provides the logic to resolve queries on actions.
+
+The exact definition of an `external entity` is left opaque, to allow for more
+interesting handlers.
+
+(TODO: Add connection to existing handlers when they are added)
+
+## Debug Action
+
+A `debug action` is essentially a marker for a type of action that may be
+performed within the compiler. There are no constraints on the granularity of an
+“action”, it can be as simple as “perform this fold” and as complex as “run this
+pass pipeline”. An action is comprised of the following:
+
+*   Tag:
+
+    -   A unique string identifier, similar to a command line flag or
+        DEBUG_TYPE.
+
+*   Description:
+
+    -   A short description of what the action represents.
+
+*   Parameter Types:
+
+    -   The types of values that are passed to queries related to this action,
+        to help guide decisions.
+
+Below is an example action that may be provided by the
+[pattern rewriter](PatternRewriter.md) framework to control the application of
+rewrite patterns.
+
+```c++
+/// A debug action that allows for controlling the application of patterns.
+/// A new action type can be defined by inheriting from `DebugAction`.
+/// * The Tag is specified via a static `StringRef getTag()` method.
+/// * The Description is specified via a static `StringRef getDescription()`
+///   method.
+/// * The parameters for the action are provided via template parameters when
+///   inheriting from `DebugAction`.
+struct ApplyPatternAction
+    : public DebugAction<Operation *, const Pattern &> {
+  static StringRef getTag() { return "apply-pattern"; }
+  static StringRef getDescription() {
+    return "Control the application of rewrite patterns";
+  }
+};
+```
+
+## Debug Action Manager
+
+The `DebugActionManager` orchestrates the various different queries relating to
+debug actions, and is accessible via the `MLIRContext`. These queries allow for
+external entities to control various aspects of the compiler. The set of
+available queries are shown below:
+
+```c++
+class DebugActionManager {
+public:
+  /// Returns true if the given action type should be executed, false otherwise.
+  /// `Params` correspond to any action specific parameters that may be used to
+  /// guide the decision.
+  template <typename ActionType, typename... Params>
+  bool shouldExecute(Params &&... params);
+};
+```
+
+Building on the example from the [previous section](#debug-action), the
+following query may be used:
+
+```c++
+/// A debug action that allows for controlling the application of patterns.
+struct ApplyPatternAction
+    : public DebugAction<Operation *, const Pattern &> {
+  static StringRef getTag() { return "apply-pattern"; }
+  static StringRef getDescription() {
+    return "Control the application of rewrite patterns";
+  }
+};
+
+// ...
+
+bool shouldApplyPattern(Operation *currentOp, const Pattern &currentPattern) {
+  MLIRContext *context = currentOp->getContext();
+  DebugActionManager &manager = context->getDebugActionManager();
+
+  // Query the action manager to see if `currentPattern` should be applied to
+  // `currentOp`.
+  return manager.shouldExecute<ApplyPatternAction>(currentOp, currentPattern);
+}
+```
+
+## Debug Action Handler
+
+A debug action handler provides the internal implementation for the various
+action related queries within the [`DebugActionManager`](debug-action-manager).
+Action handlers allow for external entities to control and inject external
+information into the compiler. Handlers can be registered with the
+`DebugActionManager` using `registerActionHandler`. There are two types of
+handlers; action-specific handlers and generic handlers.
+
+### Action Specific Handlers
+
+Action specific handlers handle a specific debug action type, with the
+parameters to its query methods mapping 1-1 to the parameter types of the action
+type. An action specific handler can be defined by inheriting from the handler
+base class defined at `ActionType::Handler` where `ActionType` is the specific
+action that should be handled. An example using our running pattern example is
+shown below:
+
+```c++
+struct MyPatternHandler : public ApplyPatternAction::Handler {
+  /// A variant of `shouldExecute` shown in the `DebugActionManager` class
+  /// above.
+  /// This method returns a FailureOr<bool>, where failure signifies that the
+  /// action was not handled (allowing for other handlers to process it), or the
+  /// boolean true/false signifying if the action should execute or not.
+  FailureOr<bool> shouldExecute(Operation *op,
+                                const RewritePattern &pattern) final;
+};
+```
+
+### Generic Handlers
+
+A generic handler allows for handling queries on any action type. These types of
+handlers are useful for implementing general functionality that doesn’t
+necessarily need to interpret the exact action parameters, or can rely on an
+external interpreter (such as the user). As these handlers are generic, they
+take a set of opaque parameters that try to map the context of the action type
+in a generic way. A generic handler can be defined by inheriting from
+`DebugActionManager::GenericHandler`. An example is shown below:
+
+```c++
+struct MyPatternHandler : public DebugActionManager::GenericHandler {
+  /// The return type of this method is the same as the action-specific handler.
+  /// The arguments to this method map the concepts of an action type in an
+  /// opaque way. The arguments are provided in such a way so that the context
+  /// of the action is still somewhat user readable, or at least loggable as
+  /// such.
+  /// - actionTag: The tag specified by the action type.
+  /// - actionDesc: The description specified by the action type.
+  /// - actionParameters: A range of StringRef that is populated by formatting
+  ///                     the parameters provided by the action into a human
+  ///                     readable form.
+  virtual FailureOr<bool> shouldExecute(StringRef actionTag,
+                                        StringRef actionDesc,
+                                        ParameterRange actionParameters);
+};
+```

mlir/include/mlir/IR/MLIRContext.h

 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
 #include <functional>
 #include <memory>
 #include <vector>
 
 namespace mlir {
 class AbstractOperation;
+class DebugActionManager;
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
 class StorageUniquer;
 
public:
   /// Returns the storage uniquer used for constructing type storage instances.
   /// This should not be used directly.
   StorageUniquer &getTypeUniquer();
 
   /// Returns the storage uniquer used for constructing attribute storage
   /// instances. This should not be used directly.
   StorageUniquer &getAttributeUniquer();
 
+  /// Returns the manager of debug actions within the context.
+  DebugActionManager &getDebugActionManager();
+
   /// These APIs are tracking whether the context will be used in a
   /// multithreading environment: this has no effect other than enabling
   /// assertions on misuses of some APIs.
   void enterMultiThreadedExecution();
   void exitMultiThreadedExecution();
 
   /// Get a dialect for the provided namespace and TypeID: abort the program if
   /// a dialect exist for this namespace with different TypeID. If a dialect has

mlir/include/mlir/Support/DebugAction.h

+//===- DebugAction.h - Debug Action Support ---------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains defintions for the debug action framework. This framework
+// allows for external entites to control certain actions taken by the compiler
+// by registering handler functions. A debug action handler provides the
+// internal implementation for the various queries on a debug action, such as
+// whether it should execute or not.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_DEBUGACTION_H
+#define MLIR_SUPPORT_DEBUGACTION_H
+
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Support/TypeID.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/TypeName.h"
+#include "llvm/Support/raw_ostream.h"
+#include <functional>
+
+namespace mlir {
+
+//===----------------------------------------------------------------------===//
+// DebugActionManager
+//===----------------------------------------------------------------------===//
+
+/// This class represents manages debug actions, and orchestrates the
+/// communication between action queries and action handlers. An action handler
+/// is either an action specific handler, i.e. a derived class of
+/// `MyActionType::Handler`, or a generic handler, i.e. a derived class of
+/// `DebugActionManager::GenericHandler`. For more details on action specific
+/// handlers, see the definition of `DebugAction::Handler` below. For more
+/// details on generic handlers, see `DebugActionManager::GenericHandler` below.
+class DebugActionManager {
+public:
+  //===--------------------------------------------------------------------===//
+  // Handlers
+  //===--------------------------------------------------------------------===//
+
+  /// This class represents the base class of a debug action handler.
+  class HandlerBase {
+  public:
+    virtual ~HandlerBase() {}
+
+    /// Return the unique handler id of this handler, use for casting
+    /// functionality.
+    TypeID getHandlerID() const { return handlerID; }
+
+  protected:
+    HandlerBase(TypeID handlerID) : handlerID(handlerID) {}
+
+    /// The type of the derived handler class. This allows for detecting if a
+    /// handler can handle a given action type.
+    TypeID handlerID;
+  };
+
+  /// This class represents a generic action handler. A generic handler allows
+  /// for handling any action type. Handlers of this type are useful for
+  /// implementing general functionality that doesn’t necessarily need to
+  /// interpret the exact action parameters, or can rely on an external
+  /// interpreter (such as the user). Given that these handlers are generic,
+  /// they take a set of opaque parameters that try to map the context of the
+  /// action type in a generic way.
+  class GenericHandler : public HandlerBase {
+  public:
+    GenericHandler() : HandlerBase(TypeID::get<GenericHandler>()) {}
+
+    /// This class represents an iterator of an opaque parameter to a generic
+    /// handler. The element of this iterator is a `StringRef` that represents a
+    /// human readable format for a given action parameter.
+    class ParameterIterator final
+        : public llvm::mapped_iterator<
+              llvm::detail::value_sequence_iterator<size_t>,
+              function_ref<StringRef(size_t)>> {
+    public:
+      using reference = StringRef;
+
+    private:
+      using llvm::mapped_iterator<
+          llvm::detail::value_sequence_iterator<size_t>,
+          function_ref<StringRef(size_t)>>::mapped_iterator;
+
+      /// Allow access to the constructor.
+      friend class DebugActionManager;
+    };
+
+    /// This class represents an range of an opaque parameters to a generic
+    /// handler. The elements of this iterator are of `StringRef` type, which
+    /// represents a human readable format for a given action parameter.
+    class ParameterRange final
+        : public llvm::iterator_range<ParameterIterator> {
+      /// Construct a parameter range from an accessor function and the length
+      /// for the number of parameters in the range.
+      ParameterRange(function_ref<StringRef(size_t)> getParamFn, size_t len)
+          : llvm::iterator_range<ParameterIterator>(
+                ParameterIterator(0, getParamFn),
+                ParameterIterator(len, getParamFn)) {}
+
+      /// Allow access to the constructor.
+      friend DebugActionManager;
+    };
+
+    /// This hook allows for controlling whether an action should execute or
+    /// not. `parameters` correspond to an opaque set of values provided by the
+    /// action as context. It should return failure if the handler could not
+    /// process the action, passing it to the next registered handler.
+    virtual FailureOr<bool> shouldExecute(StringRef actionTag,
+                                          StringRef description,
+                                          ParameterRange parameters) {
+      return failure();
+    }
+
+    /// Provide classof to allow casting between handler types.
+    static bool classof(const DebugActionManager::HandlerBase *handler) {
+      return handler->getHandlerID() == TypeID::get<GenericHandler>();
+    }
+  };
+
+  /// Register the given action handler with the manager.
+  void registerActionHandler(std::unique_ptr<HandlerBase> handler) {
+    actionHandlers.emplace_back(std::move(handler));
+  }
+  template <typename T>
+  void registerActionHandler() {
+    registerActionHandler(std::make_unique<T>());
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Action Queries
+  //===--------------------------------------------------------------------===//
+
+  /// Returns true if the given action type should be executed, false otherwise.
+  /// `Args` are a set of parameters used by handlers of `ActionType` to
+  /// determine if the action should be executed.
+  template <typename ActionType, typename... Args>
+  bool shouldExecute(Args &&...args) {
+    // Invoke the `shouldExecute` method on the provided handler.
+    auto shouldExecuteFn = [&](auto *handler, auto &&...handlerParams) {
+      return handler->shouldExecute(
+          std::forward<decltype(handlerParams)>(handlerParams)...);
+    };
+    FailureOr<bool> result = dispatchToHandler<ActionType, bool>(
+        shouldExecuteFn, std::forward<Args>(args)...);
+
+    // If the action wasn't handled, execute the action by default.
+    return succeeded(result) ? *result : true;
+  }
+
+private:
+  /// Returns true if the action manager is enabled.
+  bool isEnabled() const {
+#ifdef NDEBUG
+    // The manager is always disabled if built without debug.
+    return false;
+#else
+    return true;
+#endif
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Query to Handler Dispatch
+  //===--------------------------------------------------------------------===//
+
+  /// Dispath a given callback on any handlers that are able to process queries
+  /// on the given action type. This method returns failure if no handlers could
+  /// process the action, or success(with a result) if a handler processed the
+  /// action.
+  template <typename ActionType, typename ResultT, typename HandlerCallbackT,
+            typename... Args>
+  FailureOr<ResultT> dispatchToHandler(HandlerCallbackT &&handlerCallback,
+                                       Args &&...args) {
+    static_assert(ActionType::template canHandleWith<Args...>(),
+                  "cannot execute action with the given set of parameters");
+    if (!isEnabled())
+      return failure();
+
+    // A functor used when trying to access an opaque parameter. This is only
+    // used for GenericHandlers, and enables lazy format evaluation of opaque
+    // parameters.
+    SmallVector<std::string, sizeof...(Args)> opaqueParameters;
+    auto opaqueParamGetFn = [&](size_t index) -> StringRef {
+      // If this parameters were already formatted, access it directly.
+      if (!opaqueParameters.empty())
+        return opaqueParameters[index];
+
+      // Otherwise, format the action parameters.
+      (void)std::initializer_list<int>{
+          0, (opaqueParameters.push_back(
+                  formatParameter(std::forward<Args>(args))),
+              0)...};
+      return opaqueParameters[index];
+    };
+    GenericHandler::ParameterRange opaqueParams(opaqueParamGetFn,
+                                                sizeof...(Args));
+
+    // Process any generic or action specific handlers.
+    for (std::unique_ptr<HandlerBase> &it : llvm::reverse(actionHandlers)) {
+      FailureOr<ResultT> result = failure();
+      if (auto *handler = dyn_cast<typename ActionType::Handler>(&*it)) {
+        result = handlerCallback(handler, std::forward<Args>(args)...);
+      } else if (auto *genericHandler = dyn_cast<GenericHandler>(&*it)) {
+        result = handlerCallback(genericHandler, ActionType::getTag(),
+                                 ActionType::getDescription(), opaqueParams);
+      }
+
+      // If the handler succeeded, return the result. Otherwise, try a new
+      // handler.
+      if (succeeded(result))
+        return result;
+    }
+    return failure();
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Opaque Parameter Formatting
+  //===--------------------------------------------------------------------===//
+
+  /// Trait to check whether T can be streamed into a raw_ostream.
+  template <typename T>
+  using can_format_with_stream =
+      decltype(std::declval<llvm::raw_ostream &>() << std::declval<T &>());
+
+  /// Format the given parameter into a string containing `type: value`.
+  template <typename ParamT>
+  std::string formatParameter(ParamT &&param) {
+    std::string result;
+    llvm::raw_string_ostream os(result);
+    os << "`" << llvm::getTypeName<ParamT>() << "`";
+    formatParameterValue(param, os);
+    return os.str();
+  }
+
+  /// Format the given parameter using a raw_ostream.
+  template <typename ParamT>
+  std::enable_if_t<llvm::is_detected<can_format_with_stream, ParamT>::value>
+  formatParameterValue(ParamT &&param, llvm::raw_ostream &os) {
+    os << ": {\n  " << param << "\n}";
+  }
+  /// If there is no ostream overload, don't format anything.
+  template <typename ParamT>
+  std::enable_if_t<!llvm::is_detected<can_format_with_stream, ParamT>::value>
+  formatParameterValue(ParamT &param, llvm::raw_ostream &os) {}
+  template <typename ParamT>
+  void formatParameterValue(ParamT *param, llvm::raw_ostream &os) {
+    if (param && llvm::is_detected<can_format_with_stream, ParamT>::value)
+      formatParameterValue(*param, os);
+    else
+      os << ": {\n  " << param << "\n}";
+  }
+
+  /// The set of action parameters that have been registered with the manager.
+  SmallVector<std::unique_ptr<HandlerBase>> actionHandlers;
+};
+
+//===----------------------------------------------------------------------===//
+// DebugAction
+//===----------------------------------------------------------------------===//
+
+/// A debug action is a specific action that is to be taken by the compiler,
+/// that can be toggled and controlled by an external user. There are no
+/// constraints on the granulity of an action, it could be as simple as
+/// "perform this fold" and as complex as "run this pass pipeline". Via template
+/// parameters `ParameterTs`, a user may provide the set of argument types that
+/// are provided when handling a query on this action. Derived classes are
+/// expected to provide the following:
+///   * static llvm::StringRef getTag()
+///     - This method returns a unique string identifier, similar to a command
+///       line flag or DEBUG_TYPE.
+///   * static llvm::StringRef getDescription()
+///     - This method returns a short description of what the action represents.
+///
+/// This class provides a handler class that can be derived from to handle
+/// instances of this action. The parameters to its query methods map 1-1 to the
+/// types on the action type.
+template <typename... ParameterTs>
+class DebugAction {
+public:
+  class Handler : public DebugActionManager::HandlerBase {
+  public:
+    Handler() : HandlerBase(TypeID::get<Handler>()) {}
+
+    /// This hook allows for controlling whether an action should execute or
+    /// not. `parameters` correspond to the set of values provided by the
+    /// action as context. It should return failure if the handler could not
+    /// process the action, passing it to the next registered handler.
+    virtual FailureOr<bool> shouldExecute(ParameterTs... parameters) {
+      return failure();
+    }
+
+    /// Provide classof to allow casting between handler types.
+    static bool classof(const DebugActionManager::HandlerBase *handler) {
+      return handler->getHandlerID() ==
+             TypeID::get<DebugAction<ParameterTs...>::Handler>();
+    }
+  };
+
+private:
+  /// Returns true if the action can be handled within the given set of
+  /// parameter types.
+  template <typename... CallerParameterTs>
+  static constexpr bool canHandleWith() {
+    return llvm::is_invocable<function_ref<void(ParameterTs...)>,
+                              CallerParameterTs...>::value;
+  }
+
+  /// Allow access to `canHandleWith`.
+  friend class DebugActionManager;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_SUPPORT_DEBUGACTION_H

mlir/lib/IR/MLIRContext.cpp

 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Support/DebugAction.h"
 #include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"

 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 /// This is the implementation of the MLIRContext class, using the pImpl idiom.
 /// This class is completely private to this file, so everything is public.
 class MLIRContextImpl {
 public:
   //===--------------------------------------------------------------------===//
+  // Debugging
+  //===--------------------------------------------------------------------===//
+
+  /// An action manager for use within the context.
+  DebugActionManager debugActionManager;
+
+  //===--------------------------------------------------------------------===//
   // Identifier uniquing
   //===--------------------------------------------------------------------===//
 
   // Identifier allocator and mutex for thread safety.
   llvm::BumpPtrAllocator identifierAllocator;
   llvm::sys::SmartRWMutex<true> identifierMutex;
 
   //===--------------------------------------------------------------------===//

 static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,
                                     ArrayRef<T> elements) {
   auto result = allocator.Allocate<T>(elements.size());
   std::uninitialized_copy(elements.begin(), elements.end(), result);
   return ArrayRef<T>(result, elements.size());
 }
 
 //===----------------------------------------------------------------------===//
+// Debugging
+//===----------------------------------------------------------------------===//
+
+DebugActionManager &MLIRContext::getDebugActionManager() {
+  return getImpl().debugActionManager;
+}
+
+//===----------------------------------------------------------------------===//
 // Diagnostic Handlers
 //===----------------------------------------------------------------------===//
 
 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration

mlir/unittests/Support/CMakeLists.txt

 add_mlir_unittest(MLIRSupportTests
+  DebugActionTest.cpp
   IndentedOstreamTest.cpp
   MathExtrasTest.cpp
 )
 
 target_link_libraries(MLIRSupportTests
   PRIVATE MLIRSupportIndentedOstream MLIRSupport)

mlir/unittests/Support/DebugActionTest.cpp

+//===- DebugActionTest.cpp - Debug Action Tests ---------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/DebugAction.h"
+#include "gmock/gmock.h"
+
+// DebugActionManager is only enabled in DEBUG mode.
+#ifndef NDEBUG
+
+using namespace mlir;
+
+namespace {
+struct SimpleAction : public DebugAction<> {
+  static StringRef getTag() { return "simple-action"; }
+  static StringRef getDescription() { return "simple-action-description"; }
+};
+struct ParametricAction : public DebugAction<bool> {
+  static StringRef getTag() { return "param-action"; }
+  static StringRef getDescription() { return "param-action-description"; }
+};
+
+TEST(DebugActionTest, GenericHandler) {
+  DebugActionManager manager;
+
+  // A generic handler that always executes the simple action, but not the
+  // parametric action.
+  struct GenericHandler : public DebugActionManager::GenericHandler {
+    FailureOr<bool> shouldExecute(StringRef tag, StringRef desc,
+                                  ParameterRange parameters) final {
+      if (tag == SimpleAction::getTag()) {
+        EXPECT_EQ(desc, SimpleAction::getDescription());
+        EXPECT_EQ(llvm::size(parameters), 0);
+        return true;
+      }
+
+      EXPECT_EQ(tag, ParametricAction::getTag());
+      EXPECT_EQ(desc, ParametricAction::getDescription());
+      EXPECT_EQ(llvm::size(parameters), 1);
+      return false;
+    }
+  };
+  manager.registerActionHandler<GenericHandler>();
+
+  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
+  EXPECT_FALSE(manager.shouldExecute<ParametricAction>(true));
+}
+
+TEST(DebugActionTest, ActionSpecificHandler) {
+  DebugActionManager manager;
+
+  // Handler that simply uses the input as the decider.
+  struct ActionSpecificHandler : public ParametricAction::Handler {
+    FailureOr<bool> shouldExecute(bool shouldExecuteParam) final {
+      return shouldExecuteParam;
+    }
+  };
+  manager.registerActionHandler<ActionSpecificHandler>();
+
+  EXPECT_TRUE(manager.shouldExecute<ParametricAction>(true));
+  EXPECT_FALSE(manager.shouldExecute<ParametricAction>(false));
+
+  // There is no handler for the simple action, so it is always executed.
+  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
+}
+
+} // namespace
+
+#endif