Change the DebugAction paradigm to delegate the control to the handler

[mehdi_amini]

Summary

At the moment, we invoke shouldExecute() that way:

if (manager.shouldExecute<DebugAction>(currentOp) {
  // apply a transformation
  …
}

In this sequence, the manager isn’t involved in the actual execution
of the action and can’t develop rich instrumentations. Instead the API
could let the control to the handler itself:

// Execute the action under the control of the manager
manager.execute<DebugAction>(currentOp, [&]() {
  // apply the transformation in this callback
  …
});

This inversion of control (by injecting a callback) allows handlers to
implement potentially new interesting features: for example, snapshot
the IR before and after the action, or record an action execution time.
More importantly, it will allow to capture the nesting execution of
actions.

On the other side: handlers receives now a DebugAction object that wraps
generic information (tag and description especially) as well as
action-specific data.

Finally, the DebugActionManager is now enabled in release builds as
well.

[zero9178]

mlir/include/mlir/Support/DebugAction.h

153–154

You forgot to remove this #endif here

[mehdi_amini]

Fix the include guard

[mehdi_amini]

rebase

[rriddle]

mlir/include/mlir/Support/DebugAction.h

106

Can we drop the llvm:: on all of these?

[scotttodd]

mlir/include/mlir/Support/DebugAction.h

173–186

This had a comment before:

// TODO: We currently just pick the first handler that gives us a result,
// but in the future we may want to employ a reduction over all of the
// values returned.

Would that still be useful? It seems like being able to register multiple read-only listeners/observers would be useful, but this looks like only one can run at a time?

[jpienaar]

mlir/include/mlir/Support/DebugAction.h

93

Why the change out of debug-only for DebugAction?

148

Could you add a comment here? (Failed feels weird for me here.)

261

New editor defaults off ?

[mehdi_amini]

address comments

mlir/include/mlir/Support/DebugAction.h

93

It's in the name: this isn't "DebugAction" anymore, but Action :)

(it's broken out in various patches, but look at the follow up in D144809

148

Sure, but note that this all goes away in D144811

173–186

You won't be able to do it on the MLIRContext, this will be exposed by the ExecutionContext though. We couldn't compose things properly at this level here.

See https://reviews.llvm.org/D144812

Or the diagram in the RFC: https://discourse.llvm.org/t/rfc-mlir-action-tracing-and-debugging-mlir-based-compilers/68679/6

https://global.discourse-cdn.com/business4/uploads/llvm/original/3X/b/7/b73cd3ef38ba613acfd23305ac47ed4ddb856bbc.jpeg

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 definitions for the debug action framework. This framework
 // allows for external entities 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.
+// by registering handler functions.
 //
 //===----------------------------------------------------------------------===//
 
 #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>
 #include <type_traits>
 
 namespace mlir {
 
+/// This class represents the base class of a debug action.
+class DebugActionBase {
+public:
+  virtual ~DebugActionBase() = 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; }
+
+  virtual void print(raw_ostream &os) const {
+    os << "Action \"" << tag << "\" : " << desc << "\n";
+  }
+
+protected:
+  DebugActionBase(TypeID actionID, StringRef tag, StringRef desc)
+      : actionID(actionID), tag(tag), desc(desc) {}
+
+  /// The type of the derived action class. This allows for detecting the
+  /// specific handler of a given action type.
+  TypeID actionID;
+  StringRef tag;
+  StringRef desc;
+};
+
 //===----------------------------------------------------------------------===//
 // 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
public:
   /// 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 hook allows for controlling whether an action should execute or
-    /// not. 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) {
+    /// This hook allows for controlling the execution of an action. It should
+    /// return failure if the handler could not process the action, or whether
+    /// the `transform` was executed or not.
+    virtual FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                                    const DebugActionBase &action) {
       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) {
-    // The manager is always disabled if built without debug.
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
     actionHandlers.emplace_back(std::move(handler));
-#endif
   }
   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.
+  /// Dispatch an action represented by the `transform` callback. If no handler
+  /// is found, the `transform` callback is invoked directly.
+  /// Return true if the action was executed, false otherwise.
   template <typename ActionType, typename... Args>
-  bool shouldExecute(Args &&...args) {
-    // The manager is always disabled if built without debug.
-#if !LLVM_ENABLE_ABI_BREAKING_CHECKS
+  bool execute(llvm::function_ref<void()> transform, Args &&...args) {
+    if (actionHandlers.empty()) {
+      transform();
       return true;
-#else
-    // Invoke the `shouldExecute` method on the provided handler.
-    auto shouldExecuteFn = [&](auto *handler, auto &&...handlerParams) {
-      return handler->shouldExecute(
-          std::forward<decltype(handlerParams)>(handlerParams)...);
+    }
+
+    // Invoke the `execute` method on the provided handler.
+    auto executeFn = [&](auto *handler, auto &&...handlerParams) {
+      return handler->execute(
+          transform,
+          ActionType(std::forward<decltype(handlerParams)>(handlerParams)...));
     };
     FailureOr<bool> result = dispatchToHandler<ActionType, bool>(
-        shouldExecuteFn, std::forward<Args>(args)...);
+        executeFn, std::forward<Args>(args)...);
+    if (failed(result)) {
+      transform();
+      return true;
+    }
 
-    // If the action wasn't handled, execute the action by default.
-    return succeeded(result) ? *result : true;
-#endif
+    // Return the result of the handler.
+    return *result;
   }
 
 private:
-// The manager is always disabled if built without debug.
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
   //===--------------------------------------------------------------------===//
   // 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");
 
     // Process any generic or action specific handlers.
-    // TODO: We currently just pick the first handler that gives us a result,
-    // but in the future we may want to employ a reduction over all of the
-    // values returned.
+    // The first handler that gives us a result is the one that we will return.
     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());
+        result = handlerCallback(genericHandler, std::forward<Args>(args)...);
       }
 
       // If the handler succeeded, return the result. Otherwise, try a new
       // handler.
       if (succeeded(result))
         return result;
     }
     return failure();
   }
 
   /// The set of action handlers that have been registered with the manager.
   SmallVector<std::unique_ptr<HandlerBase>> actionHandlers;
-#endif
 };
 
 //===----------------------------------------------------------------------===//
 // 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()
 ///     - 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 {
+class DebugAction : public DebugActionBase {
 public:
+  DebugAction()
+      : DebugActionBase(TypeID::get<Derived>(), Derived::getTag(),
+                        Derived::getDescription()) {}
+
+  /// Provide classof to allow casting between action types.
+  static bool classof(const DebugActionBase *action) {
+    return action->getActionID() == TypeID::get<Derived>();
+  }
+
   class Handler : public DebugActionManager::HandlerBase {
   public:
     Handler() : HandlerBase(TypeID::get<Derived>()) {}
 
-    /// This hook allows for controlling whether an action should execute or
-    /// not. `parameters` correspond to the set of values provided by the
+    /// This hook allows for controlling the execution of an action.
+    /// `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) {
+    virtual FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                                    const Derived &action) {
       return failure();
     }
 
     /// Provide classof to allow casting between handler types.
     static bool classof(const DebugActionManager::HandlerBase *handler) {
       return handler->getHandlerID() == TypeID::get<Derived>();
     }
   };
 
 private:
   /// Returns true if the action can be handled within the given set of
   /// parameter types.
   template <typename... CallerParameterTs>
   static constexpr bool canHandleWith() {
     return std::is_invocable_v<function_ref<void(ParameterTs...)>,
                                CallerParameterTs...>;
   }
 
   /// Allow access to `canHandleWith`.
   friend class DebugActionManager;
 };
 
 } // namespace mlir
 
-#endif // MLIR_SUPPORT_DEBUGACTION_H
+#endif // MLIR_SUPPORT_DEBUGACTION_H
 No newline at end of file

mlir/include/mlir/Support/DebugCounter.h

public:
   /// Add a counter for the given debug 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.
   void addCounter(StringRef actionTag, int64_t countToSkip,
                   int64_t countToStopAfter);
 
   /// Register a counter with the specified name.
-  FailureOr<bool> shouldExecute(StringRef tag, StringRef description) final;
+  FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                          const DebugActionBase &action) final;
 
   /// Print the counters that have been registered with this instance to the
   /// provided output stream.
   void print(raw_ostream &os) const;
 
   /// Register the command line options for debug counters.
   static void registerCLOptions();
 

mlir/lib/Support/DebugCounter.cpp

 void DebugCounter::addCounter(StringRef actionTag, int64_t countToSkip,
                               int64_t countToStopAfter) {
   assert(!counters.count(actionTag) &&
          "a counter for the given action was already registered");
   counters.try_emplace(actionTag, countToSkip, countToStopAfter);
 }
 
 // Register a counter with the specified name.
-FailureOr<bool> DebugCounter::shouldExecute(StringRef tag,
-                                            StringRef description) {
-  auto counterIt = counters.find(tag);
+FailureOr<bool> DebugCounter::execute(llvm::function_ref<void()> transform,
+                                      const DebugActionBase &action) {
+  auto counterIt = counters.find(action.getTag());
   if (counterIt == counters.end())
     return true;
 
   ++counterIt->second.count;
 
   // We only execute while the `countToSkip` is not smaller than `count`, and
   // `countToStopAfter + countToSkip` is larger than `count`. Negative counters
   // always execute.

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 "mlir/Support/TypeID.h"
 #include "gmock/gmock.h"
 
-// DebugActionManager is only enabled in DEBUG mode.
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
-
 using namespace mlir;
 
 namespace {
 struct SimpleAction : DebugAction<SimpleAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SimpleAction)
   static StringRef getTag() { return "simple-action"; }
   static StringRef getDescription() { return "simple-action-description"; }
 };
 struct OtherSimpleAction : DebugAction<OtherSimpleAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OtherSimpleAction)
   static StringRef getTag() { return "other-simple-action"; }
   static StringRef getDescription() {
     return "other-simple-action-description";
   }
 };
 struct ParametricAction : DebugAction<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"; }
 };
 
 TEST(DebugActionTest, GenericHandler) {
   DebugActionManager manager;
 
   // A generic handler that always executes the simple action, but not the
   // parametric action.
   struct GenericHandler : DebugActionManager::GenericHandler {
-    FailureOr<bool> shouldExecute(StringRef tag, StringRef desc) final {
-      if (tag == SimpleAction::getTag()) {
+    FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                            const DebugActionBase &action) final {
+      StringRef desc = action.getDescription();
+      if (isa<SimpleAction>(action)) {
         EXPECT_EQ(desc, SimpleAction::getDescription());
+        transform();
         return true;
       }
 
-      EXPECT_EQ(tag, ParametricAction::getTag());
+      EXPECT_TRUE(isa<ParametricAction>(action));
       EXPECT_EQ(desc, ParametricAction::getDescription());
       return false;
     }
   };
   manager.registerActionHandler<GenericHandler>();
 
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
-  EXPECT_FALSE(manager.shouldExecute<ParametricAction>(true));
+  auto noOp = []() { return; };
+  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
+  EXPECT_FALSE(manager.execute<ParametricAction>(noOp, true));
 }
 
 TEST(DebugActionTest, ActionSpecificHandler) {
   DebugActionManager manager;
 
   // Handler that simply uses the input as the decider.
   struct ActionSpecificHandler : ParametricAction::Handler {
-    FailureOr<bool> shouldExecute(bool shouldExecuteParam) final {
-      return shouldExecuteParam;
+    FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                            const ParametricAction &action) final {
+      if (action.executeParam)
+        transform();
+      return action.executeParam;
     }
   };
   manager.registerActionHandler<ActionSpecificHandler>();
 
-  EXPECT_TRUE(manager.shouldExecute<ParametricAction>(true));
-  EXPECT_FALSE(manager.shouldExecute<ParametricAction>(false));
+  int count = 0;
+  auto incCount = [&]() { count++; };
+  EXPECT_TRUE(manager.execute<ParametricAction>(incCount, true));
+  EXPECT_EQ(count, 1);
+  EXPECT_FALSE(manager.execute<ParametricAction>(incCount, false));
+  EXPECT_EQ(count, 1);
 
   // There is no handler for the simple action, so it is always executed.
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
+  EXPECT_TRUE(manager.execute<SimpleAction>(incCount));
+  EXPECT_EQ(count, 2);
 }
 
 TEST(DebugActionTest, DebugCounterHandler) {
   DebugActionManager manager;
 
   // Handler that uses the number of action executions as the decider.
   struct DebugCounterHandler : SimpleAction::Handler {
-    FailureOr<bool> shouldExecute() final { return numExecutions++ < 3; }
+    FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                            const SimpleAction &action) final {
+      bool shouldExecute = numExecutions++ < 3;
+      if (shouldExecute)
+        transform();
+      return shouldExecute;
+    }
     unsigned numExecutions = 0;
   };
   manager.registerActionHandler<DebugCounterHandler>();
 
   // Check that the action is executed 3 times, but no more after.
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
-  EXPECT_FALSE(manager.shouldExecute<SimpleAction>());
-  EXPECT_FALSE(manager.shouldExecute<SimpleAction>());
+  auto noOp = []() { return; };
+  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
+  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
+  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
+  EXPECT_FALSE(manager.execute<SimpleAction>(noOp));
+  EXPECT_FALSE(manager.execute<SimpleAction>(noOp));
 }
 
 TEST(DebugActionTest, NonOverlappingActionSpecificHandlers) {
   DebugActionManager manager;
 
   // One handler returns true and another returns false
   struct SimpleActionHandler : SimpleAction::Handler {
-    FailureOr<bool> shouldExecute() final { return true; }
+    FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                            const SimpleAction &action) final {
+      transform();
+      return true;
+    }
   };
   struct OtherSimpleActionHandler : OtherSimpleAction::Handler {
-    FailureOr<bool> shouldExecute() final { return false; }
+    FailureOr<bool> execute(llvm::function_ref<void()> transform,
+                            const OtherSimpleAction &action) final {
+      transform();
+      return false;
+    }
   };
   manager.registerActionHandler<SimpleActionHandler>();
   manager.registerActionHandler<OtherSimpleActionHandler>();
-  EXPECT_TRUE(manager.shouldExecute<SimpleAction>());
-  EXPECT_FALSE(manager.shouldExecute<OtherSimpleAction>());
+  auto noOp = []() { return; };
+  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
+  EXPECT_FALSE(manager.execute<OtherSimpleAction>(noOp));
 }
 
 } // namespace
-
-#endif

mlir/unittests/Support/DebugCounterTest.cpp

 //===- DebugCounterTest.cpp - Debug Counter 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/DebugCounter.h"
 #include "mlir/Support/TypeID.h"
 #include "gmock/gmock.h"
 
 using namespace mlir;
 
-// DebugActionManager is only enabled in DEBUG mode.
-#if LLVM_ENABLE_ABI_BREAKING_CHECKS
-
 namespace {
 
 struct CounterAction : public DebugAction<CounterAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterAction)
   static StringRef getTag() { return "counter-action"; }
   static StringRef getDescription() { return "Test action for debug counters"; }
 };
 
 TEST(DebugCounterTest, CounterTest) {
   std::unique_ptr<DebugCounter> counter = std::make_unique<DebugCounter>();
   counter->addCounter(CounterAction::getTag(), /*countToSkip=*/1,
                       /*countToStopAfter=*/3);
 
   DebugActionManager manager;
   manager.registerActionHandler(std::move(counter));
 
+  auto noOp = []() { return; };
+
   // The first execution is skipped.
-  EXPECT_FALSE(manager.shouldExecute<CounterAction>());
+  EXPECT_FALSE(manager.execute<CounterAction>(noOp));
 
   // The counter stops after 3 successful executions.
-  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
-  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
-  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
-  EXPECT_FALSE(manager.shouldExecute<CounterAction>());
+  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
+  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
+  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
+  EXPECT_FALSE(manager.execute<CounterAction>(noOp));
 }
 
 } // namespace
-
-#endif