Rename DebugAction to tracing::Action and move related code from lib/Support to lib/IR and lib/Debug

This is a preparation for adding support for more infrastructure around the concept
of Action and make tracing Action more of a first class concept.
The doc will be updated later in a subsequent revision after the changes are
completed.

Action belongs to IR because of circular dependency: Actions are dispatched through
the MLIRContext but Action will learn to encapsulate IR construct.

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

mlir/include/mlir/Debug/Counter.h

@@ -6,31 +6,32 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_SUPPORT_DEBUGCOUNTER_H
-#define MLIR_SUPPORT_DEBUGCOUNTER_H
+#ifndef MLIR_TRACING_DEBUGCOUNTER_H
+#define MLIR_TRACING_DEBUGCOUNTER_H
 
-#include "mlir/Support/DebugAction.h"
+#include "mlir/IR/Action.h"
 #include "llvm/ADT/StringMap.h"
 #include <string>
 
 namespace mlir {
+namespace tracing {
 
-/// This class implements a debug action handler that attaches a counter value
+/// This class implements an action handler that attaches a counter value
 /// to debug actions and enables/disables execution of these action based on the
 /// value of the counter. The counter controls the execution of the action with
 /// a "skip" and "count" value. The "skip" value is used to skip a certain
-/// number of initial executions of a debug action. The "count" value is used to
-/// prevent a debug action from executing after it has executed for a set number
+/// number of initial executions of an action. The "count" value is used to
+/// prevent an action from executing after it has executed for a set number
 /// of times (not including any executions that have been skipped). For example,
-/// a counter for a debug action with `skip=47` and `count=2`, would skip the
+/// a counter for an action with `skip=47` and `count=2`, would skip the
 /// first 47 executions, then execute twice, and finally prevent any further
 /// executions.
-class DebugCounter : public DebugActionManager::GenericHandler {
+class DebugCounter : public ActionManager::GenericHandler {
 public:
   DebugCounter();
   ~DebugCounter() override;
 
-  /// Add a counter for the given debug action tag. `countToSkip` is the number
+  /// Add a counter for the given action tag. `countToSkip` is the number
   /// of counter executions to skip before enabling execution of the action.
   /// `countToStopAfter` is the number of executions of the counter to allow
   /// before preventing the action from executing any more.
@@ -39,7 +40,7 @@ public:
 
   /// Register a counter with the specified name.
   FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                          const DebugActionBase &action) final;
+                          const Action &action) final;
 
   /// Print the counters that have been registered with this instance to the
   /// provided output stream.
@@ -69,6 +70,7 @@ private:
   llvm::StringMap<Counter> counters;
 };
 
+} // namespace tracing
 } // namespace mlir
 
-#endif
+#endif // MLIR_TRACING_DEBUGCOUNTER_H

mlir/include/mlir/IR/Action.h

@@ -1,4 +1,4 @@
-//===- DebugAction.h - Debug Action Support ---------------------*- C++ -*-===//
+//===- Action.h -  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.
@@ -6,14 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains definitions for the debug action framework. This framework
+// This file contains definitions for the action framework. This framework
 // allows for external entities to control certain actions taken by the compiler
 // by registering handler functions.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_SUPPORT_DEBUGACTION_H
-#define MLIR_SUPPORT_DEBUGACTION_H
+#ifndef MLIR_IR_ACTION_H
+#define MLIR_IR_ACTION_H
 
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/TypeID.h"
@@ -26,53 +26,58 @@
 #include <type_traits>
 
 namespace mlir {
+namespace tracing {
 
-/// This class represents the base class of a debug action.
-class DebugActionBase {
+/// An 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 granularity 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.
+///
+/// This class represents the base class of the ActionImpl class (see below).
+/// This holds the template-invariant elements of the Action class.
+class Action {
 public:
-  virtual ~DebugActionBase() = default;
+  virtual ~Action() = default;
 
   /// Return the unique action id of this action, use for casting
   /// functionality.
   TypeID getActionID() const { return actionID; }
 
-  StringRef getTag() const { return tag; }
-
-  StringRef getDescription() const { return desc; }
+  /// Return a string "tag" which intends to uniquely identify this type of
+  /// action. For example "pass-application" or "pattern-rewrite".
+  virtual StringRef getTag() const = 0;
 
   virtual void print(raw_ostream &os) const {
-    os << "Action \"" << tag << "\" : " << desc << "\n";
+    os << "Action \"" << getTag() << "\"";
   }
 
 protected:
-  DebugActionBase(TypeID actionID, StringRef tag, StringRef desc)
-      : actionID(actionID), tag(tag), desc(desc) {}
+  Action(TypeID actionID) : actionID(actionID) {}
 
-  /// The type of the derived action class. This allows for detecting the
-  /// specific handler of a given action type.
+  /// The type of the derived action class, used for `isa`/`dyn_cast`.
   TypeID actionID;
-  StringRef tag;
-  StringRef desc;
 };
 
 //===----------------------------------------------------------------------===//
-// DebugActionManager
+// ActionManager
 //===----------------------------------------------------------------------===//
 
-/// This class represents manages debug actions, and orchestrates the
+/// This class represents manages 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 {
+/// `ActionManager::GenericHandler`. For more details on action specific
+/// handlers, see the definition of `Action::Handler` below. For more
+/// details on generic handlers, see `ActionManager::GenericHandler` below.
+class ActionManager {
 public:
   //===--------------------------------------------------------------------===//
   // Handlers
   //===--------------------------------------------------------------------===//
 
-  /// This class represents the base class of a debug action handler.
+  /// This class represents the base class of an action handler.
   class HandlerBase {
   public:
     virtual ~HandlerBase() = default;
@@ -104,12 +109,12 @@ public:
     /// return failure if the handler could not process the action, or whether
     /// the `transform` was executed or not.
     virtual FailureOr<bool> execute(function_ref<void()> transform,
-                                    const DebugActionBase &action) {
+                                    const Action &action) {
       return failure();
     }
 
     /// Provide classof to allow casting between handler types.
-    static bool classof(const DebugActionManager::HandlerBase *handler) {
+    static bool classof(const ActionManager::HandlerBase *handler) {
       return handler->getHandlerID() == TypeID::get<GenericHandler>();
     }
   };
@@ -193,39 +198,26 @@ private:
   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 granularity 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()
+/// CRTP Implementation of an action. This class provides a base class for
+/// implementing specific actions.
+///  Derived classes are expected to provide the following:
+///   * static constexpr StringLiteral tag = "...";
 ///     - 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 Derived, typename... ParameterTs>
-class DebugAction : public DebugActionBase {
+class ActionImpl : public Action {
 public:
-  DebugAction()
-      : DebugActionBase(TypeID::get<Derived>(), Derived::getTag(),
-                        Derived::getDescription()) {}
+  ActionImpl() : Action(TypeID::get<Derived>()) {}
 
   /// Provide classof to allow casting between action types.
-  static bool classof(const DebugActionBase *action) {
+  static bool classof(const Action *action) {
     return action->getActionID() == TypeID::get<Derived>();
   }
 
-  class Handler : public DebugActionManager::HandlerBase {
+  /// Forward tag access to the derived class.
+  StringRef getTag() const final { return Derived::tag; }
+
+  class Handler : public ActionManager::HandlerBase {
   public:
     Handler() : HandlerBase(TypeID::get<Derived>()) {}
 
@@ -239,7 +231,7 @@ public:
     }
 
     /// Provide classof to allow casting between handler types.
-    static bool classof(const DebugActionManager::HandlerBase *handler) {
+    static bool classof(const ActionManager::HandlerBase *handler) {
       return handler->getHandlerID() == TypeID::get<Derived>();
     }
   };
@@ -254,9 +246,10 @@ private:
   }
 
   /// Allow access to `canHandleWith`.
-  friend class DebugActionManager;
+  friend class ActionManager;
 };
 
+} // namespace tracing
 } // namespace mlir
 
-#endif // MLIR_SUPPORT_DEBUGACTION_H
+#endif // MLIR_IR_ACTION_H

mlir/include/mlir/IR/MLIRContext.h

@@ -20,7 +20,9 @@ class ThreadPool;
 } // namespace llvm
 
 namespace mlir {
-class DebugActionManager;
+namespace tracing {
+class ActionManager;
+}
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
@@ -30,6 +32,7 @@ class Location;
 class MLIRContextImpl;
 class RegisteredOperationName;
 class StorageUniquer;
+class IRUnit;
 
 /// MLIRContext is the top-level object for a collection of MLIR operations. It
 /// holds immortal uniqued objects like types, and the tables used to unique
@@ -215,7 +218,7 @@ public:
   StorageUniquer &getAttributeUniquer();
 
   /// Returns the manager of debug actions within the context.
-  DebugActionManager &getDebugActionManager();
+  tracing::ActionManager &getActionManager();
 
   /// These APIs are tracking whether the context will be used in a
   /// multithreading environment: this has no effect other than enabling

mlir/lib/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(Analysis)
 add_subdirectory(AsmParser)
 add_subdirectory(Bytecode)
 add_subdirectory(Conversion)
+add_subdirectory(Debug)
 add_subdirectory(Dialect)
 add_subdirectory(IR)
 add_subdirectory(Interfaces)

mlir/lib/Debug/CMakeLists.txt

@@ -0,0 +1,11 @@
+add_mlir_library(MLIRDebug
+  DebugCounter.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Debug
+
+  LINK_LIBS PUBLIC
+  ${LLVM_PTHREAD_LIB}
+  MLIRIR
+  MLIRSupport
+)

mlir/lib/Debug/DebugCounter.cpp

@@ -6,13 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Support/DebugCounter.h"
+#include "mlir/Debug/Counter.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/ManagedStatic.h"
 
 using namespace mlir;
+using namespace mlir::tracing;
 
 //===----------------------------------------------------------------------===//
 // DebugCounter CommandLine Options
@@ -63,7 +64,7 @@ void DebugCounter::addCounter(StringRef actionTag, int64_t countToSkip,
 
 // Register a counter with the specified name.
 FailureOr<bool> DebugCounter::execute(llvm::function_ref<void()> transform,
-                                      const DebugActionBase &action) {
+                                      const Action &action) {
   auto counterIt = counters.find(action.getTag());
   if (counterIt == counters.end())
     return true;

mlir/lib/IR/MLIRContext.cpp

@@ -12,6 +12,7 @@
 #include "AttributeDetail.h"
 #include "IntegerSetDetail.h"
 #include "TypeDetail.h"
+#include "mlir/IR/Action.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
@@ -23,7 +24,6 @@
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Types.h"
-#include "mlir/Support/DebugAction.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallString.h"
@@ -124,7 +124,7 @@ public:
   //===--------------------------------------------------------------------===//
 
   /// An action manager for use within the context.
-  DebugActionManager debugActionManager;
+  tracing::ActionManager actionManager;
 
   //===--------------------------------------------------------------------===//
   // Diagnostics
@@ -348,8 +348,8 @@ static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,
 // Debugging
 //===----------------------------------------------------------------------===//
 
-DebugActionManager &MLIRContext::getDebugActionManager() {
-  return getImpl().debugActionManager;
+tracing::ActionManager &MLIRContext::getActionManager() {
+  return getImpl().actionManager;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Support/CMakeLists.txt

@@ -1,5 +1,4 @@
 set(LLVM_OPTIONAL_SOURCES
-  DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   InterfaceSupport.cpp
@@ -10,7 +9,6 @@ set(LLVM_OPTIONAL_SOURCES
 )
 
 add_mlir_library(MLIRSupport
-  DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   InterfaceSupport.cpp

mlir/lib/Tools/mlir-opt/CMakeLists.txt

@@ -6,6 +6,7 @@ add_mlir_library(MLIROptLib
 
   LINK_LIBS PUBLIC
   MLIRBytecodeWriter
+  MLIRDebug
   MLIRPass
   MLIRParser
   MLIRSupport

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

@@ -13,6 +13,7 @@
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/Debug/Counter.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -23,7 +24,6 @@
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
-#include "mlir/Support/DebugCounter.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/Timing.h"
 #include "mlir/Support/ToolUtilities.h"
@@ -212,7 +212,7 @@ static LogicalResult processBuffer(raw_ostream &os,
   context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
   if (config.shouldVerifyDiagnostics())
     context.printOpOnDiagnostic(false);
-  context.getDebugActionManager().registerActionHandler<DebugCounter>();
+  context.getActionManager().registerActionHandler<tracing::DebugCounter>();
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
@@ -324,7 +324,7 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   registerDefaultTimingManagerCLOptions();
-  DebugCounter::registerCLOptions();
+  tracing::DebugCounter::registerCLOptions();
 
   // Build the list of dialects as a header for the --help message.
   std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();

mlir/unittests/CMakeLists.txt

@@ -6,6 +6,7 @@ endfunction()
 
 add_subdirectory(Analysis)
 add_subdirectory(Conversion)
+add_subdirectory(Debug)
 add_subdirectory(Dialect)
 add_subdirectory(Interfaces)
 add_subdirectory(IR)

mlir/unittests/Debug/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_mlir_unittest(MLIRDebugTests
+  DebugCounterTest.cpp
+)
+
+target_link_libraries(MLIRDebugTests
+  PRIVATE MLIRDebug)

mlir/unittests/Debug/DebugCounterTest.cpp

@@ -6,26 +6,26 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Support/DebugCounter.h"
+#include "mlir/Debug/Counter.h"
 #include "mlir/Support/TypeID.h"
 #include "gmock/gmock.h"
 
 using namespace mlir;
+using namespace mlir::tracing;
 
 namespace {
 
-struct CounterAction : public DebugAction<CounterAction> {
+struct CounterAction : public ActionImpl<CounterAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterAction)
-  static StringRef getTag() { return "counter-action"; }
-  static StringRef getDescription() { return "Test action for debug counters"; }
+  static constexpr StringLiteral tag = "counter-action";
 };
 
 TEST(DebugCounterTest, CounterTest) {
   std::unique_ptr<DebugCounter> counter = std::make_unique<DebugCounter>();
-  counter->addCounter(CounterAction::getTag(), /*countToSkip=*/1,
+  counter->addCounter(CounterAction::tag, /*countToSkip=*/1,
                       /*countToStopAfter=*/3);
 
-  DebugActionManager manager;
+  ActionManager manager;
   manager.registerActionHandler(std::move(counter));
 
   auto noOp = []() { return; };

mlir/unittests/IR/ActionTest.cpp

@@ -1,4 +1,4 @@
-//===- DebugActionTest.cpp - Debug Action Tests ---------------------------===//
+//===- ActionTest.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.
@@ -6,50 +6,45 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Support/DebugAction.h"
+#include "mlir/IR/Action.h"
 #include "mlir/Support/TypeID.h"
 #include "gmock/gmock.h"
 
 using namespace mlir;
+using namespace mlir::tracing;
 
 namespace {
-struct SimpleAction : DebugAction<SimpleAction> {
+struct SimpleAction : ActionImpl<SimpleAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SimpleAction)
-  static StringRef getTag() { return "simple-action"; }
-  static StringRef getDescription() { return "simple-action-description"; }
+  static constexpr StringLiteral tag = "simple-action";
 };
-struct OtherSimpleAction : DebugAction<OtherSimpleAction> {
+struct OtherSimpleAction : ActionImpl<OtherSimpleAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OtherSimpleAction)
-  static StringRef getTag() { return "other-simple-action"; }
-  static StringRef getDescription() {
-    return "other-simple-action-description";
-  }
+  static constexpr StringLiteral tag = "other-simple-action";
 };
-struct ParametricAction : DebugAction<ParametricAction, bool> {
+struct ParametricAction : ActionImpl<ParametricAction, bool> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ParametricAction)
   ParametricAction(bool executeParam) : executeParam(executeParam) {}
   bool executeParam;
-  static StringRef getTag() { return "param-action"; }
-  static StringRef getDescription() { return "param-action-description"; }
+  static constexpr StringLiteral tag = "param-action";
 };
 
-TEST(DebugActionTest, GenericHandler) {
-  DebugActionManager manager;
+TEST(ActionTest, GenericHandler) {
+  ActionManager manager;
 
   // A generic handler that always executes the simple action, but not the
   // parametric action.
-  struct GenericHandler : DebugActionManager::GenericHandler {
+  struct GenericHandler : ActionManager::GenericHandler {
     FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                            const DebugActionBase &action) final {
-      StringRef desc = action.getDescription();
+                            const Action &action) final {
+      StringRef tag = action.getTag();
       if (isa<SimpleAction>(action)) {
-        EXPECT_EQ(desc, SimpleAction::getDescription());
+        EXPECT_EQ(tag, SimpleAction::tag);
         transform();
         return true;
       }
 
       EXPECT_TRUE(isa<ParametricAction>(action));
-      EXPECT_EQ(desc, ParametricAction::getDescription());
       return false;
     }
   };
@@ -60,8 +55,8 @@ TEST(DebugActionTest, GenericHandler) {
   EXPECT_FALSE(manager.execute<ParametricAction>(noOp, true));
 }
 
-TEST(DebugActionTest, ActionSpecificHandler) {
-  DebugActionManager manager;
+TEST(ActionTest, ActionSpecificHandler) {
+  ActionManager manager;
 
   // Handler that simply uses the input as the decider.
   struct ActionSpecificHandler : ParametricAction::Handler {
@@ -86,8 +81,8 @@ TEST(DebugActionTest, ActionSpecificHandler) {
   EXPECT_EQ(count, 2);
 }
 
-TEST(DebugActionTest, DebugCounterHandler) {
-  DebugActionManager manager;
+TEST(ActionTest, DebugCounterHandler) {
+  ActionManager manager;
 
   // Handler that uses the number of action executions as the decider.
   struct DebugCounterHandler : SimpleAction::Handler {
@@ -111,8 +106,8 @@ TEST(DebugActionTest, DebugCounterHandler) {
   EXPECT_FALSE(manager.execute<SimpleAction>(noOp));
 }
 
-TEST(DebugActionTest, NonOverlappingActionSpecificHandlers) {
-  DebugActionManager manager;
+TEST(ActionTest, NonOverlappingActionSpecificHandlers) {
+  ActionManager manager;
 
   // One handler returns true and another returns false
   struct SimpleActionHandler : SimpleAction::Handler {

mlir/unittests/IR/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_unittest(MLIRIRTests
+  ActionTest.cpp
   AdaptorTest.cpp
   AttributeTest.cpp
   DialectTest.cpp

mlir/unittests/Support/CMakeLists.txt

@@ -1,6 +1,4 @@
 add_mlir_unittest(MLIRSupportTests
-  DebugActionTest.cpp
-  DebugCounterTest.cpp
   IndentedOstreamTest.cpp
   MathExtrasTest.cpp
   StorageUniquerTest.cpp