diff --git a/mlir/docs/DebugActions.md b/mlir/docs/DebugActions.md new file mode 100644 --- /dev/null +++ b/mlir/docs/DebugActions.md @@ -0,0 +1,171 @@ +# 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`](#debug-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`](#debug-action-manager) that will + provide an answer as to what behavior the action should take. +* An external entity registers an [`action handler`](#debug-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. The set of possible action queries is detailed in the +[`action manager`](#debug-action-manager) section below. + +(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 { + 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 via +[action handlers](#debug-action-handler). When resolving a query for an action, +the result from the most recently registered handler is used. + +TODO: It may be interesting to support merging results from multiple action +handlers, but this is left for future work when driven by a real use case. + +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 + bool shouldExecute(Params &&... params); +}; +``` + +Building on the example from the [previous section](#debug-action), an example +usage of the `shouldExecute` query is shown below: + +```c++ +/// A debug action that allows for controlling the application of patterns. +struct ApplyPatternAction + : public DebugAction { + static StringRef getTag() { return "apply-pattern"; } + static StringRef getDescription() { + return "Control the application of rewrite patterns"; + } +}; + +// ... + +bool shouldApplyPattern(Operation *currentOp, const Pattern ¤tPattern) { + MLIRContext *context = currentOp->getContext(); + DebugActionManager &manager = context->getDebugActionManager(); + + // Query the action manager to see if `currentPattern` should be applied to + // `currentOp`. + return manager.shouldExecute(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, 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 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. + virtual FailureOr shouldExecute(StringRef actionTag, + StringRef actionDesc); +}; +``` diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h --- a/mlir/include/mlir/IR/MLIRContext.h +++ b/mlir/include/mlir/IR/MLIRContext.h @@ -17,6 +17,7 @@ namespace mlir { class AbstractOperation; +class DebugActionManager; class DiagnosticEngine; class Dialect; class DialectRegistry; @@ -150,6 +151,9 @@ /// 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. diff --git a/mlir/include/mlir/Support/DebugAction.h b/mlir/include/mlir/Support/DebugAction.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Support/DebugAction.h @@ -0,0 +1,229 @@ +//===- 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 + +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()) {} + + /// 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 shouldExecute(StringRef actionTag, + StringRef description) { + return failure(); + } + + /// Provide classof to allow casting between handler types. + static bool classof(const DebugActionManager::HandlerBase *handler) { + return handler->getHandlerID() == TypeID::get(); + } + }; + + /// Register the given action handler with the manager. + void registerActionHandler(std::unique_ptr handler) { + // The manager is always disabled if built without debug. +#ifndef NDEBUG + actionHandlers.emplace_back(std::move(handler)); +#endif + } + template + void registerActionHandler() { + registerActionHandler(std::make_unique()); + } + + //===--------------------------------------------------------------------===// + // 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 + bool shouldExecute(Args &&...args) { + // The manager is always disabled if built without debug. +#ifdef NDEBUG + return true; +#else + // Invoke the `shouldExecute` method on the provided handler. + auto shouldExecuteFn = [&](auto *handler, auto &&...handlerParams) { + return handler->shouldExecute( + std::forward(handlerParams)...); + }; + FailureOr result = dispatchToHandler( + shouldExecuteFn, std::forward(args)...); + + // If the action wasn't handled, execute the action by default. + return succeeded(result) ? *result : true; +#endif + } + +private: +// The manager is always disabled if built without debug. +#ifndef NDEBUG + //===--------------------------------------------------------------------===// + // 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 + FailureOr dispatchToHandler(HandlerCallbackT &&handlerCallback, + Args &&...args) { + static_assert(ActionType::template canHandleWith(), + "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. + for (std::unique_ptr &it : llvm::reverse(actionHandlers)) { + FailureOr result = failure(); + if (auto *handler = dyn_cast(&*it)) { + result = handlerCallback(handler, std::forward(args)...); + } else if (auto *genericHandler = dyn_cast(&*it)) { + result = handlerCallback(genericHandler, ActionType::getTag(), + ActionType::getDescription()); + } + + // 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> 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 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 +class DebugAction { +public: + class Handler : public DebugActionManager::HandlerBase { + public: + Handler() : HandlerBase(TypeID::get()) {} + + /// 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 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::Handler>(); + } + }; + +private: + /// Returns true if the action can be handled within the given set of + /// parameter types. + template + static constexpr bool canHandleWith() { + return llvm::is_invocable, + CallerParameterTs...>::value; + } + + /// Allow access to `canHandleWith`. + friend class DebugActionManager; +}; + +} // end namespace mlir + +#endif // MLIR_SUPPORT_DEBUGACTION_H diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -24,6 +24,7 @@ #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" @@ -210,6 +211,13 @@ /// 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 //===--------------------------------------------------------------------===// @@ -412,6 +420,14 @@ return ArrayRef(result, elements.size()); } +//===----------------------------------------------------------------------===// +// Debugging +//===----------------------------------------------------------------------===// + +DebugActionManager &MLIRContext::getDebugActionManager() { + return getImpl().debugActionManager; +} + //===----------------------------------------------------------------------===// // Diagnostic Handlers //===----------------------------------------------------------------------===// diff --git a/mlir/unittests/Support/CMakeLists.txt b/mlir/unittests/Support/CMakeLists.txt --- a/mlir/unittests/Support/CMakeLists.txt +++ b/mlir/unittests/Support/CMakeLists.txt @@ -1,4 +1,5 @@ add_mlir_unittest(MLIRSupportTests + DebugActionTest.cpp IndentedOstreamTest.cpp MathExtrasTest.cpp ) diff --git a/mlir/unittests/Support/DebugActionTest.cpp b/mlir/unittests/Support/DebugActionTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Support/DebugActionTest.cpp @@ -0,0 +1,88 @@ +//===- 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 { + 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 shouldExecute(StringRef tag, StringRef desc) final { + if (tag == SimpleAction::getTag()) { + EXPECT_EQ(desc, SimpleAction::getDescription()); + return true; + } + + EXPECT_EQ(tag, ParametricAction::getTag()); + EXPECT_EQ(desc, ParametricAction::getDescription()); + return false; + } + }; + manager.registerActionHandler(); + + EXPECT_TRUE(manager.shouldExecute()); + EXPECT_FALSE(manager.shouldExecute(true)); +} + +TEST(DebugActionTest, ActionSpecificHandler) { + DebugActionManager manager; + + // Handler that simply uses the input as the decider. + struct ActionSpecificHandler : public ParametricAction::Handler { + FailureOr shouldExecute(bool shouldExecuteParam) final { + return shouldExecuteParam; + } + }; + manager.registerActionHandler(); + + EXPECT_TRUE(manager.shouldExecute(true)); + EXPECT_FALSE(manager.shouldExecute(false)); + + // There is no handler for the simple action, so it is always executed. + EXPECT_TRUE(manager.shouldExecute()); +} + +TEST(DebugActionTest, DebugCounterHandler) { + DebugActionManager manager; + + // Handler that uses the number of action executions as the decider. + struct DebugCounterHandler : public SimpleAction::Handler { + FailureOr shouldExecute() final { return numExecutions++ < 3; } + unsigned numExecutions = 0; + }; + manager.registerActionHandler(); + + // Check that the action is executed 3 times, but no more after. + EXPECT_TRUE(manager.shouldExecute()); + EXPECT_TRUE(manager.shouldExecute()); + EXPECT_TRUE(manager.shouldExecute()); + EXPECT_FALSE(manager.shouldExecute()); + EXPECT_FALSE(manager.shouldExecute()); +} + +} // namespace + +#endif