Delete ActionManager and replace it with a simple callback on the Context

[mehdi_amini]

Summary

The concept of the ActionManager acts as a sort of "Hub" that can receive
various types of action and dispatch them to a set of registered handlers.
One handler will handle the action or it'll cascade to other handlers.

This model does not really fit the current evolution of the Action tracing
and debugging: we can't foresee a good case where this behavior compose with
the use-case behind the handlers. Instead we simplify it with a single
callback installed on the Context.

[mehdi_amini]

rebase

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

231–245

The naming here feels much too general for something specific to actions, I'd rather have this be an explicitly named API.

Also, why "dispatch"? instead of something like "execute" or something more aligned with how actions define the verbage.

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

210

[mehdi_amini]

mlir/include/mlir/IR/MLIRContext.h

231–245

I am fine with execute, I went back and forth on it and ended up using dispatch() to encourage feedback :)

What made me hesitate is that an Action can be skipped (by debug counters for example).

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

231–245

I'd be fine with dispatch, but really just want it to be consistent. I think we use execute in the action verbiage (unless that got fixed?). Mainly just want it to be clear what the API is for when reading it, and be able to attribute it to actions.

[mehdi_amini]

mlir/include/mlir/IR/MLIRContext.h

231–245

Could use delegate() as well (dispatch may feel like we may use a thread pool or something like that).

You also said it is too generic, do you mean you would prefer to see something like context.delegateAction() ?

The call sites will look like context.execute<PassAction>() right now, that would look like context.executeAction<PassAction>() instead.

I haven't written the doc for all this, I'll add a revision when we're done with all the revisions (so that I document the final state).

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

231–245

Yeah, I'm a bit concerned about the API being very generically named. I was less concerned about the (templated) callsites, and more with users reading API (needs a lot of context to see what's going on). Maybe this concern is overblown. Can you wrap the API in a comment block? Might help with grouping/reading:

//===--------------------------------------------------------------------===//
// Action API
//===--------------------------------------------------------------------===//

/// Dispatch the provided action to the handler if any, or just execute it.
void dispatch(function_ref<void()> actionFn, const tracing::Action &action) {
  if (LLVM_UNLIKELY(hasActionHandler()))
    dispatchInternal(actionFn, action);
  else
    actionFn();
}

[jpienaar]

mlir/include/mlir/IR/MLIRContext.h

231–245

No I think that's valid concern, I'd even suggest consisting affixing (ActionHandlerTy, executeAction)

[mehdi_amini]

Rename the API as executeAction()

mlir/include/mlir/Debug/Counter.h

 /// 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 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 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 ActionManager::GenericHandler {
+class DebugCounter {
 public:
   DebugCounter();
-  ~DebugCounter() override;
+  ~DebugCounter();
 
   /// 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.
   void addCounter(StringRef actionTag, int64_t countToSkip,
                   int64_t countToStopAfter);
 
-  /// Register a counter with the specified name.
-  FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                          const Action &action) final;
+  /// Entry point for handling actions.
+  void operator()(llvm::function_ref<void()> transform, const Action &action);
 
   /// 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();
 
 private:
+  // Returns true if the next action matching this tag should be executed.
+  bool shouldExecute(StringRef tag);
+
   /// Apply the registered CL options to this debug counter instance.
   void applyCLOptions();
 
   /// This struct represents a specific counter being tracked.
   struct Counter {
     Counter(int64_t countToSkip = 0, int64_t countToStopAfter = -1)
         : countToSkip(countToSkip), countToStopAfter(countToStopAfter) {}
 

mlir/include/mlir/IR/Action.h

 class Operation;
 class Region;
 class Value;
 namespace tracing {
 
 /// 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.
+/// "perform this fold" and as complex as "run this pass pipeline".
 ///
 /// 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 ~Action() = default;
 
   /// Return the unique action id of this action, use for casting

 
 protected:
   Action(TypeID actionID) : actionID(actionID) {}
 
   /// The type of the derived action class, used for `isa`/`dyn_cast`.
   TypeID actionID;
 };
 
-//===----------------------------------------------------------------------===//
-// ActionManager
-//===----------------------------------------------------------------------===//
-
-/// 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
-/// `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 an action handler.
-  class HandlerBase {
-  public:
-    virtual ~HandlerBase() = default;
-
-    /// 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 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 Action &action) {
-      return failure();
-    }
-
-    /// Provide classof to allow casting between handler types.
-    static bool classof(const ActionManager::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
-  //===--------------------------------------------------------------------===//
-
-  /// 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 execute(llvm::function_ref<void()> transform, Args &&...args) {
-    if (actionHandlers.empty()) {
-      transform();
-      return true;
-    }
-
-    // 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>(
-        executeFn, std::forward<Args>(args)...);
-    if (failed(result)) {
-      transform();
-      return true;
-    }
-
-    // Return the result of the handler.
-    return *result;
-  }
-
-private:
-  //===--------------------------------------------------------------------===//
-  // 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.
-    // 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, 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;
-};
-
 /// 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.
-template <typename Derived, typename... ParameterTs>
+template <typename Derived>
 class ActionImpl : public Action {
 public:
   ActionImpl() : Action(TypeID::get<Derived>()) {}
 
   /// Provide classof to allow casting between action types.
   static bool classof(const Action *action) {
     return action->getActionID() == TypeID::get<Derived>();
   }
 
   /// 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>()) {}
-
-    /// 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> execute(llvm::function_ref<void()> transform,
-                                    const Derived &action) {
-      return failure();
-    }
-
-    /// Provide classof to allow casting between handler types.
-    static bool classof(const ActionManager::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 ActionManager;
 };
 
 } // namespace tracing
 } // namespace mlir
 
 #endif // MLIR_IR_ACTION_H

mlir/include/mlir/IR/MLIRContext.h

 #include <vector>
 
 namespace llvm {
 class ThreadPool;
 } // namespace llvm
 
 namespace mlir {
 namespace tracing {
-class ActionManager;
+class Action;
 }
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
 class DynamicDialect;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
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.
-  tracing::ActionManager &getActionManager();
-
   /// 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
   /// not been loaded for this namespace/TypeID yet, use the provided ctor to
   /// create one on the fly and load it. Returns a pointer to the dialect owned
   /// by the context.
   /// The use of this method is in general discouraged in favor of
   /// 'getOrLoadDialect<DialectClass>()'.
   Dialect *getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
                             function_ref<std::unique_ptr<Dialect>()> ctor);
 
   /// Returns a hash of the registry of the context that may be used to give
   /// a rough indicator of if the state of the context registry has changed. The
   /// context registry correlates to loaded dialects and their entities
   /// (attributes, operations, types, etc.).
   llvm::hash_code getRegistryHash();
 
+  //===--------------------------------------------------------------------===//
+  // Action API
+  //===--------------------------------------------------------------------===//
+
+  /// Signatures for the action handler that can be registered with the context.
+  using HandlerTy =
+      std::function<void(function_ref<void()>, const tracing::Action &)>;
+
+  /// Register a handler for handling actions that are dispatched through this
+  /// context. A nullptr handler can be set to disable a previously set handler.
+  void registerActionHandler(HandlerTy handler);
+
+  /// Return true if a valid ActionHandler is set.
+  bool hasActionHandler();
+
+  /// Dispatch the provided action to the handler if any, or just execute it.
+  void executeAction(function_ref<void()> actionFn,
+                     const tracing::Action &action) {
+    if (LLVM_UNLIKELY(hasActionHandler()))
+      executeActionInternal(actionFn, action);
+    else
+      actionFn();
+  }
+
+  /// Dispatch the provided action to the handler if any, or just execute it.
+  template <typename ActionTy, typename... Args>
+  void executeAction(function_ref<void()> actionFn, Args &&...args) {
+    if (LLVM_UNLIKELY(hasActionHandler()))
+      executeAction(actionFn, ActionTy(std::forward<Args>(args)...));
+    else
+      actionFn();
+  }
+
 private:
   /// Return true if the given dialect is currently loading.
   bool isDialectLoading(StringRef dialectNamespace);
 
+  /// Internal helper for the dispatch method.
+  void executeActionInternal(function_ref<void()> actionFn,
+                             const tracing::Action &action);
+
   const std::unique_ptr<MLIRContextImpl> impl;
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 
 //===----------------------------------------------------------------------===//
 // MLIRContext CommandLine Options

mlir/lib/Debug/DebugCounter.cpp

 /// before preventing the action from executing any more.
 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::execute(llvm::function_ref<void()> transform,
-                                      const Action &action) {
-  auto counterIt = counters.find(action.getTag());
+void DebugCounter::operator()(llvm::function_ref<void()> transform,
+  if (shouldExecute(action.getTag()))
+    transform();
+}
+
+bool DebugCounter::shouldExecute(StringRef tag) {
+  auto counterIt = counters.find(tag);
   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/lib/IR/MLIRContext.cpp

 #include "mlir/IR/Types.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/RWMutex.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/raw_ostream.h"
 #include <memory>
 #include <optional>
 

 /// 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.
-  tracing::ActionManager actionManager;
+  /// An action handler for handling actions that are dispatched through this
+  /// context.
+  std::function<void(function_ref<void()>, const tracing::Action &)>
+      actionHandler;
 
   //===--------------------------------------------------------------------===//
   // Diagnostics
   //===--------------------------------------------------------------------===//
   DiagnosticEngine diagEngine;
 
   //===--------------------------------------------------------------------===//
   // Options

 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
+// Action Handling
 //===----------------------------------------------------------------------===//
 
-tracing::ActionManager &MLIRContext::getActionManager() {
-  return getImpl().actionManager;
+void MLIRContext::registerActionHandler(HandlerTy handler) {
+  getImpl().actionHandler = std::move(handler);
 }
 
+/// Dispatch the provided action to the handler if any, or just execute it.
+void MLIRContext::executeActionInternal(function_ref<void()> actionFn,
+                                        const tracing::Action &action) {
+  assert(getImpl().actionHandler);
+  getImpl().actionHandler(actionFn, action);
+}
+
+bool MLIRContext::hasActionHandler() { return (bool)getImpl().actionHandler; }
+
 //===----------------------------------------------------------------------===//
 // Diagnostic Handlers
 //===----------------------------------------------------------------------===//
 
 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//

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

if (threadPool)
     context.setThreadPool(*threadPool);
 
   // Parse the input file.
   if (config.shouldPreloadDialectsInContext())
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
   if (config.shouldVerifyDiagnostics())
     context.printOpOnDiagnostic(false);
-  context.getActionManager().registerActionHandler<tracing::DebugCounter>();
+  context.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.
   if (!config.shouldVerifyDiagnostics()) {
     SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
     return performActions(os, sourceMgr, &context, config);
   }
 

mlir/unittests/Debug/DebugCounterTest.cpp

 namespace {
 
 struct CounterAction : public ActionImpl<CounterAction> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterAction)
   static constexpr StringLiteral tag = "counter-action";
 };
 
 TEST(DebugCounterTest, CounterTest) {
-  std::unique_ptr<DebugCounter> counter = std::make_unique<DebugCounter>();
-  counter->addCounter(CounterAction{}.getTag(), /*countToSkip=*/1,
-                      /*countToStopAfter=*/3);
+  DebugCounter counter;
+  counter.addCounter(CounterAction{}.getTag(), /*countToSkip=*/1,
 
-  ActionManager manager;
-  manager.registerActionHandler(std::move(counter));
-
-  auto noOp = []() { return; };
+  int count = 0;
+  auto noOp = [&]() {
+    ++count;
+    return;
+  };
 
   // The first execution is skipped.
-  EXPECT_FALSE(manager.execute<CounterAction>(noOp));
+  counter(noOp, CounterAction{});
+  EXPECT_EQ(count, 0);
 
   // The counter stops after 3 successful executions.
-  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
-  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
-  EXPECT_TRUE(manager.execute<CounterAction>(noOp));
-  EXPECT_FALSE(manager.execute<CounterAction>(noOp));
+  counter(noOp, CounterAction{});
+  EXPECT_EQ(count, 1);
+  counter(noOp, CounterAction{});
+  EXPECT_EQ(count, 2);
+  counter(noOp, CounterAction{});
+  EXPECT_EQ(count, 3);
+  counter(noOp, CounterAction{});
+  EXPECT_EQ(count, 3);
 }
 
 } // namespace

mlir/unittests/IR/ActionTest.cpp

-//===- 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.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/IR/Action.h"
-#include "mlir/Support/TypeID.h"
-#include "gmock/gmock.h"
-
-using namespace mlir;
-using namespace mlir::tracing;
-
-namespace {
-struct SimpleAction : ActionImpl<SimpleAction> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SimpleAction)
-  static constexpr StringLiteral tag = "simple-action";
-};
-struct OtherSimpleAction : ActionImpl<OtherSimpleAction> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OtherSimpleAction)
-  static constexpr StringLiteral tag = "other-simple-action";
-};
-struct ParametricAction : ActionImpl<ParametricAction, bool> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ParametricAction)
-  ParametricAction(bool executeParam) : executeParam(executeParam) {}
-  bool executeParam;
-  static constexpr StringLiteral tag = "param-action";
-};
-
-TEST(ActionTest, GenericHandler) {
-  ActionManager manager;
-
-  // A generic handler that always executes the simple action, but not the
-  // parametric action.
-  struct GenericHandler : ActionManager::GenericHandler {
-    FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                            const Action &action) final {
-      StringRef tag = action.getTag();
-      if (isa<SimpleAction>(action)) {
-        EXPECT_EQ(tag, SimpleAction{}.getTag());
-        transform();
-        return true;
-      }
-
-      EXPECT_TRUE(isa<ParametricAction>(action));
-      return false;
-    }
-  };
-  manager.registerActionHandler<GenericHandler>();
-
-  auto noOp = []() { return; };
-  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
-  EXPECT_FALSE(manager.execute<ParametricAction>(noOp, true));
-}
-
-TEST(ActionTest, ActionSpecificHandler) {
-  ActionManager manager;
-
-  // Handler that simply uses the input as the decider.
-  struct ActionSpecificHandler : ParametricAction::Handler {
-    FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                            const ParametricAction &action) final {
-      if (action.executeParam)
-        transform();
-      return action.executeParam;
-    }
-  };
-  manager.registerActionHandler<ActionSpecificHandler>();
-
-  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.execute<SimpleAction>(incCount));
-  EXPECT_EQ(count, 2);
-}
-
-TEST(ActionTest, DebugCounterHandler) {
-  ActionManager manager;
-
-  // Handler that uses the number of action executions as the decider.
-  struct DebugCounterHandler : SimpleAction::Handler {
-    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.
-  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(ActionTest, NonOverlappingActionSpecificHandlers) {
-  ActionManager manager;
-
-  // One handler returns true and another returns false
-  struct SimpleActionHandler : SimpleAction::Handler {
-    FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                            const SimpleAction &action) final {
-      transform();
-      return true;
-    }
-  };
-  struct OtherSimpleActionHandler : OtherSimpleAction::Handler {
-    FailureOr<bool> execute(llvm::function_ref<void()> transform,
-                            const OtherSimpleAction &action) final {
-      transform();
-      return false;
-    }
-  };
-  manager.registerActionHandler<SimpleActionHandler>();
-  manager.registerActionHandler<OtherSimpleActionHandler>();
-  auto noOp = []() { return; };
-  EXPECT_TRUE(manager.execute<SimpleAction>(noOp));
-  EXPECT_FALSE(manager.execute<OtherSimpleAction>(noOp));
-}
-
-} // namespace

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
-  ActionTest.cpp
   AdaptorTest.cpp
   AttributeTest.cpp
   DialectTest.cpp
   InterfaceTest.cpp
   IRMapping.cpp
   InterfaceAttachmentTest.cpp
   OperationSupportTest.cpp
   PatternMatchTest.cpp