Somehow it was escaped in these documents.
12 KiB
Action: Tracing and Debugging MLIR-based Compilers
[TOC]
See also the slides and the recording from the MLIR Open Meeting where this feature was demoed.
Overview
Action
are means to encapsulate any transformation of any granularity in a way
that can be intercepted by the framework for debugging or tracing purposes,
including skipping a transformation programmatically (think about "compiler
fuel" or "debug counters" in LLVM). As such, "executing a pass" is an Action, so
is "try to apply one canonicalization pattern", or "tile this loop".
In MLIR, passes and patterns are the main abstractions to encapsulate general IR transformations. The primary way of observing transformations along the way is to enable “debug printing” of the IR (e.g. -mlir-print-ir-after-all to print after each pass execution). On top of this, finer grain tracing may be available with -debug which enables more detailed logs from the transformations themselves. However, this method has some scaling issues: it is limited to a single stream of text that can be gigantic and requires tedious crawling through this log a posteriori. Iterating through multiple runs of collecting such logs and analyzing it can be very time consuming and often not very practical beyond small input programs.
The Action
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
Action
execution is shown below:
- Compiler developer defines an
Action
class, that is representing the transformation or utility that they are developing. - Depending on the needs, the developer identifies single unit of
transformations, and dispatch them to the
MLIRContext
for execution. - An external entity registers an "action handler" with the action manager, and provides the logic surrounding the transformation execution.
The exact definition of an external entity
is left opaque, to allow for more
interesting handlers.
Wrapping a Transformation in an Action
There are two parts for getting started with enabling tracing through Action in
existing or new code: 1) defining an actual Action
class, and 2) encapsulating
the transformation in a lambda function.
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:
/// A custom Action can be defined minimally by deriving from
/// `tracing::ActionImpl`.
class MyCustomAction : public tracing::ActionImpl<MyCustomAction> {
public:
using Base = tracing::ActionImpl<MyCustomAction>;
/// Actions are initialized with an array of IRUnit (that is either Operation,
/// Block, or Region) that provide context for the IR affected by a transformation.
MyCustomAction(ArrayRef<IRUnit> irUnits)
: Base(irUnits) {}
/// This tag should uniquely identify this action, it can be matched for filtering
/// during processing.
static constexpr StringLiteral tag = "unique-tag-for-my-action";
static constexpr StringLiteral desc =
"This action will encapsulate a some very specific transformation";
};
Any transformation can then be dispatched with this Action
through the
MLIRContext
:
context->executeAction<ApplyPatternAction>(
[&]() {
rewriter.setInsertionPoint(op);
...
},
/*IRUnits=*/{op, region});
An action can also carry arbitrary payload, for example we can extend the
MyCustomAction
class above with the following member:
/// A custom Action can be defined minimally by deriving from
/// `tracing::ActionImpl`. It can has any members!
class MyCustomAction : public tracing::ActionImpl<MyCustomAction> {
public:
using Base = tracing::ActionImpl<MyCustomAction>;
/// Actions are initialized with an array of IRUnit (that is either Operation,
/// Block, or Region) that provide context for the IR affected by a transformation.
/// Other constructor arguments can also be required here.
MyCustomAction(ArrayRef<IRUnit> irUnits, int count, PaddingStyle padding)
: Base(irUnits), count(count), padding(padding) {}
/// This tag should uniquely identify this action, it can be matched for filtering
/// during processing.
static constexpr StringLiteral tag = "unique-tag-for-my-action";
static constexpr StringLiteral desc =
"This action will encapsulate a some very specific transformation";
/// Extra members can be carried by the Action
int count;
PaddingStyle padding;
};
These new members must then be passed as arguments when dispatching an Action
:
context->executeAction<ApplyPatternAction>(
[&]() {
rewriter.setInsertionPoint(op);
...
},
/*IRUnits=*/{op, region},
/*count=*/count,
/*padding=*/padding);
Intercepting Actions
When a transformation is executed through an Action
, it can be directly
intercepted via a handler that can be set on the MLIRContext
:
/// 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);
This handler takes two arguments: the first on is the transformation wrapped in a callback, and the second is a reference to the associated action object. The handler has full control of the execution, as such it can also decide to return without executing the callback, skipping the transformation entirely!
MLIR-provided Handlers
MLIR provides some predefined action handlers for immediate use that are believed to be useful for most projects built with MLIR.
Debug Counters
When debugging a compiler issue,
"bisection"
is a useful technique for locating the root cause of the issue. Debug Counters
enable using this technique for debug actions by attaching a counter value to a
specific action and enabling/disabling execution of this 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 of times (not
including any executions that have been skipped). If the "skip" value is
negative, the action will always execute. If the "count" value is negative, the
action will always execute after the "skip" value has been reached. For example,
a counter for a debug action with skip=47
and count=2
, would skip the first
47 executions, then execute twice, and finally prevent any further executions.
With a bit of tooling, the values to use for the counter can be automatically
selected; allowing for finding the exact execution of a debug action that
potentially causes the bug being investigated.
Note: The DebugCounter action handler does not support multi-threaded execution,
and should only be used in MLIRContexts where multi-threading is disabled (e.g.
via -mlir-disable-threading
).
CommandLine Configuration
The DebugCounter
handler provides several that allow for configuring counters.
The main option is mlir-debug-counter
, which accepts a comma separated list of
<count-name>=<counter-value>
. A <counter-name>
is the debug action tag to
attach the counter, suffixed with either -skip
or -count
. A -skip
suffix
will set the "skip" value of the counter. A -count
suffix will set the "count"
value of the counter. The <counter-value>
component is a numeric value to use
for the counter. An example is shown below using MyCustomAction
defined above:
$ mlir-opt foo.mlir -mlir-debug-counter=unique-tag-for-my-action-skip=47,unique-tag-for-my-action-count=2
The above configuration would skip the first 47 executions of
ApplyPatternAction
, then execute twice, and finally prevent any further
executions.
Note: Each counter currently only has one skip
and one count
value, meaning
that sequences of skip
/count
will not be chained.
The mlir-print-debug-counter
option may be used to print out debug counter
information after all counters have been accumulated. The information is printed
in the following format:
DebugCounter counters:
<action-tag> : {<current-count>,<skip>,<count>}
For example, using the options above we can see how many times an action is executed:
$ mlir-opt foo.mlir -mlir-debug-counter=unique-tag-for-my-action-skip=-1 -mlir-print-debug-counter --pass-pipeline="builtin.module(func.func(my-pass))" --mlir-disable-threading
DebugCounter counters:
unique-tag-for-my-action : {370,-1,-1}
ExecutionContext
The ExecutionContext
is a component that provides facility to unify the kind
of functionalities that most compiler debuggers tool would need, exposed in a
composable way.
The ExecutionContext
is itself registered as a handler with the MLIRContext
and tracks all executed actions, keeping a per-thread stack of action execution.
It acts as a middleware that handles the flow of action execution while allowing
injection and control from a debugger.
- Multiple
Observers
can be registered with theExecutionContext
. When an action is dispatched for execution, it is passed to each of theObservers
before and after executing the transformation. - Multiple
BreakpointManager
can be registered with theExecutionContext
. When an action is dispatched for execution, it is passed to each of the registeredBreakpointManager
until one matches the action and return a validBreakpoint
object. In this case, the "callback" set by the client on theExecutionContext
is invoked, otherwise the transformation is directly executed. - A single callback:
using CallbackTy = function_ref<Control(const ActionActiveStack *)>;
can be registered with theExecutionContext
, it is invoked when aBreakPoint
is hit by anAction
. The returned value of typeControl
is an enum instructing theExecutionContext
of how to proceed next:
Since the callback actually controls the execution, there can be only one registered at any given time./// Enum that allows the client of the context to control the execution of the /// action. /// - Apply: The action is executed. /// - Skip: The action is skipped. /// - Step: The action is executed and the execution is paused before the next /// action, including for nested actions encountered before the /// current action finishes. /// - Next: The action is executed and the execution is paused after the /// current action finishes before the next action. /// - Finish: The action is executed and the execution is paused only when we /// reach the parent/enclosing operation. If there are no enclosing /// operation, the execution continues without stopping. enum Control { Apply = 1, Skip = 2, Step = 3, Next = 4, Finish = 5 };
Debugger ExecutionContext Hook
MLIR provides a callback for the ExecutionContext
that implements a small
runtime suitable for debuggers like gdb
or lldb
to interactively control the
execution. It can be setup with
mlir::setupDebuggerExecutionContextHook(executionContext);
or using mlir-opt
with the --mlir-enable-debugger-hook
flag. This runtime exposes a set of C API
function that can be called from a debugger to:
- set breakpoints matching either action tags, or the
FileLineCol
locations of the IR associated with the action. - set the
Control
flag to be returned to theExecutionContext
. - control a "cursor" allowing to navigate through the IR and inspect it from the IR context associated with the action.
The implementation of this runtime can serve as an example for other implementation of programmatic control of the execution.
Logging Observer
One observer is provided that allows to log action execution on a provided
stream. It can be exercised with mlir-opt
using --log-actions-to=<filename>
,
and optionally filtering the output with
--log-mlir-actions-filter=<FileLineCol>
. This observer is not thread-safe at
the moment.