Transform Dialect
Fine-grain transformation control dialect
Disclaimer ¶
** Proceed with care: not ready for general use. **
This dialect is evolving rapidly and may change on a very short notice. To decrease the maintenance burden and churn, only a few in-tree use cases are currently supported in the main tree:
- high-level transformations on “structured ops” (i.e. ops that operate on chunks of data in a way that can be decomposed into operations on smaller chunks of data and control flow) in Linalg, Tensor and Vector dialects.
Please post a description of the intended use case on the MLIR forum and wait for confirmation.
Overview ¶
This dialect provides operations that can be used to control transformation of the IR using a different portion of the IR. It refers to the IR being transformed as payload IR, and to the IR guiding the transformation as transform IR.
The main use case for this dialect is orchestrating fine-grain transformations on individual operations or sets thereof. For example, it may involve finding loop-like operations with specific properties (e.g., large size) in the payload IR, applying loop tiling to those and only those operations, and then applying loop unrolling to the inner loops produced by the previous transformations. As such, it is not intended as a replacement for the pass infrastructure, nor for the pattern rewriting infrastructure. In the most common case, the transform IR will be processed and applied to the payload IR by a pass. Transformations expressed by the transform dialect may be implemented using the pattern infrastructure or any other relevant MLIR component.
The following IR gives a rough idea of what the operations in this dialect may look like:
%0 = transform.loop.find { size > 42 }
%1:2 = transform.loop.tile { tile_sizes = [2,3,4] }
transform.loop.unroll %1#1
The values defined by operations in this dialect correspond to (groups of)
operations in the payload IR. In the example above, %0
corresponds to the
set of loops found in the payload IR that satisfy the condition, and %1
correspond to groups of outer and inner loops, respectively, produced by
the tiling transformation.
Overall, Transform IR ops are expected to be contained in a single top-level
op. Such top-level ops specify how to apply the transformations described
by the operations they contain, e.g., transform.sequence
executes
transformations one by one and fails if any of them fails. Such ops are
expected to have the PossibleTopLevelTransformOpTrait
and may be used
without arguments.
Dialect Extension Mechanism ¶
This dialect is designed to be extensible, that is, clients of this dialect
are allowed to inject additional operations into this dialect using the
TransformDialectExtension
mechanism. This allows the dialect to avoid a
dependency on the implementation of the transformation as well as to avoid
introducing dialect-specific transform dialects. In the example above,
the operations may have been injected by a notional loop
dialect rather
than defined in this dialect, hence the common prefix.
It is recommended to prefix injected operations with one or several
dot-separated words that indicate which extension adds them. For
dialect-specific transformations, the prefix is naturally the name of the
dialect, e.g., transform.affine.reschedule
. For dialect-agnostic
transformations (typically implemented using interfaces), the prefix may
be derived from the interface name or from a common concept, e.g.,
transform.loop.tile
may apply to any loop-like operation that implements
TileableOpInterface
. The C++ classes for the dialect extension should
include the prefix in their name, e.g., AffineTransformDialectExtension
or
LoopTransformDialectExtension
in the cases above. Unprefixed operation
names are reserved for ops defined directly in the Transform dialect.
Operations injected into the dialect must:
Implement the
TransformOpInterface
to execute the corresponding transformation on the payload IR.Implement the
MemoryEffectsOpInterface
to annotate the effects of the transform IR operation on the payload IR as well as on the mapping between transform IR values and payload IR operations. See below for the description of available effects.
The presence of interface implementations is checked at runtime when the dialect is loaded to allow for those implementations to be supplied by separate dialect extensions if desired.
Side Effects ¶
The Transform dialect relies on MLIR side effect modelling to enable optimization of the transform IR. More specifically, it provides several side effect resource objects and expects operations to describe their effects on these resources.
TransformMappingResource
- side effect resource corresponding to the mapping between transform IR values and payload IR operations.An
Allocate
effect from this resource means creating a new mapping entry, it is always accompanied by aWrite
effet.A
Read
effect from this resource means accessing the mapping.A
Free
effect on this resource indicates the removal of the mapping entry, typically after a transformation that modifies the payload IR operations associated with one of the transform IR operation’s operands. It is always accompanied by aRead
effect.
PayloadIRResource
- side effect resource corresponding to the payload IR itself.A
Read
effect from this resource means accessing the payload IR.A
Write
effect on this resource means mutating the payload IR. It is almost always accompanied by aRead
.
The typical flow of values in the transform IR is as follows. Most
operations produce new transform IR values and immediately associate them
with a list of payload IR operations. This corresponds to Allocate
and
Write
effects on the TransformMappingResource
, and often requires at
least a Read
effect on the PayloadIRResource
. Transform operations that
only inspect the payload IR to produce new handles are usually limited to
these effects on their operands. Transform operations that mutate the
payload IR are thought to consume the handles provided as operands, that
is have the Read
and Free
effects on them. As with the usual memory
effects, using a value after it was freed is incorrect. In case of the
transform IR, this value is likely associated with payload IR operations
that were modified or even removed by the transformation, so it is
meaningless to refer to them. When further transformations are desired, the
transform operations can return new handles that can be read or consumed
by subsequent operations.
Intended Use and Integrations ¶
The transformation control infrastructure provided by this dialect is positioned roughly between rewrite patterns and passes. A transformation that is executed by a transform operation is likely to be sufficiently complex to require at least a set of patterns to be implemented. It is also expected to be more focused than a pass: a pass typically applies identical transformations everywhere in the IR, a transform dialect-controlled transformation would apply to a small subset of operations selected, e.g., by a pattern-matching operation or generated by a previous transformation. It is discouraged, although technically possible, to run a pass pipeline as part of the transform op implementation.
One of the main scenarios for using this dialect is fine-grain chaining of transformations. For example, a loop-like operation may see its iteration domain split into two parts, implemented as separate loops (transformation known as index-set splitting), each of which is then transformed differently (e.g., the first loop is tiled and the second unrolled) with the necessary enabling and cleanup patterns around the main transformation:
// <generate %loop, e.g., by pattern-matching>
// ...
%parts:2 = transform.loop.split %loop { upper_bound_divisible_by = 8 }
transform.loop.tile %parts#0 { tile_sizes = [8] }
transform.loop.unroll %parts#1 { full }
This composition would have been difficult to implement as separate passes since the hypothetical “tiling” and “unrolling” pass would need to somehow differentiate between the parts of the loop produced by the previous pass (both are the same operation, and it is likely undesirable to pollute the operation with pass-specific information). Implementing passes that run the combined transfomration would have run into the combinatorial explosion issue due to multiple possible transform compositions or into the need for deep pass parameterization, the ultimate form of which is an ad-hoc dialect to specify which transformations the pass should run. The transform dialect provides a uniform, extensible mechanism for controlling transformations in such cases.
The transform dialect is supposed to be consumed by an “interpreter” pass
that drives the application of transformations. To ensure extensibility and
composability, this pass is not expected to actually perform the
transformations specified by the ops. Instead, the transformations are
implemented by the transform ops themselves via TransformOpInterface
. The
pass serves as the entry point, handles the flow of transform operations and
takes care of bookkeeping. As such, the transform dialect does not provide
the interpreter pass. Instead, it provides a set of utilities that can be
used by clients to define their own interpreter passes or as part of a more
complex pass. For example, the mapping between values in the tranfsorm IR
and operations in the payload IR, or the function that applies the
transformations specified by ops in the given block sequentially. Note that
a transform op may have regions with further transform ops in them, with
the op itself guiding how to dispatch the transformation control flow to
those regions. This approach allows clients to decide on the relative
location of the transform IR in their input (e.g., nested modules, separate
modules, optional regions to certain operations, etc.), register additional
transform operations and perform client-specific bookkeeping.
Effects on the Infrastructure ¶
Although scoped to a single dialect, this functionality conceptually belongs to the MLIR infrastructure. It aims to be minimally intrusive and opt-in.
Some infrastructural components may grow extra functionality to support the transform dialect. In particular, the pattern infrastructure may add extra hooks to identify the “main results” of a transformation or to notify external observers about changes made to certain operations. These are not expected to affect the existing uses of the infrastructure.
For the sake of reusability, transformations should be implemented as utility functions that are called from the interface methods of transform ops rather than having the methods directly act on the payload IR.
Operation definition ¶
transform.pdl_match
(::mlir::transform::PDLMatchOp) ¶
Finds ops that match the named PDL pattern
Syntax:
operation ::= `transform.pdl_match` $pattern_name `in` $root attr-dict
Find Payload IR ops nested within the Payload IR op associated with the
operand that match the PDL pattern identified by its name. The pattern is
expected to be defined in the closest surrounding WithPDLPatternsOp
.
Produces a Transform IR value associated with the list of Payload IR ops that matched the pattern. The order of results in the list is that of the Operation::walk, clients are advised not to rely on a specific order though. If the operand is assocaited with multiple Payload IR ops, finds matching ops nested within each of those and produces a single list containing all of the matched ops.
The tranfsormation is considered successful regardless of whether some Payload IR ops actually matched the pattern and only fails if the pattern could not be looked up or compiled.
Interfaces: TransformOpInterface
Attributes: ¶
Attribute | MLIR Type | Description |
---|---|---|
pattern_name | ::mlir::SymbolRefAttr | symbol reference attribute |
Operands: ¶
Operand | Description |
---|---|
root | PDL handle to an mlir::Operation * |
Results: ¶
Result | Description |
---|---|
matched | PDL handle to an mlir::Operation * |
transform.sequence
(::mlir::transform::SequenceOp) ¶
Contains a sequence of other transform ops to apply
Syntax:
operation ::= `transform.sequence` ($root^)? attr-dict-with-keyword regions (`:` type($results)^)?
The transformations indicated by the sequence are applied in order of their appearance. Each value produced by a transformation within the sequence corresponds to an operation or a group of operations in the payload IR. During application, if any transformation in the sequence fails, the entire sequence fails immediately leaving the payload IR in potentially invalid state, i.e., this operation offers no transformation rollback capabilities.
The entry block of this operation has a single argument that maps to either the operand if provided or the top-level container operation of the payload IR, typically the root operation of the pass interpreting the transform dialect. Operand omission is only allowed for sequences not contained in another sequence.
Traits: PossibleTopLevelTransformOpTrait, SingleBlockImplicitTerminator<::mlir::transform::YieldOp>
Interfaces: MemoryEffectOpInterface, OpAsmOpInterface, TransformOpInterface
Operands: ¶
Operand | Description |
---|---|
root | PDL handle to an mlir::Operation * |
Results: ¶
Result | Description |
---|---|
results | any type |
transform.with_pdl_patterns
(::mlir::transform::WithPDLPatternsOp) ¶
Contains PDL patterns available for use in transforms
Syntax:
operation ::= `transform.with_pdl_patterns` ($root^)? attr-dict-with-keyword regions
This op contains a set of named PDL patterns that are available for the Transform dialect operations to be used for pattern matching. For example, PDLMatchOp can be used to produce a Transform IR value associated with all Payload IR operations that match the pattern as follows:
transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation):
pdl.pattern @my_pattern : benefit(1) {
%0 = pdl.operation //...
// Regular PDL goes here.
pdl.rewrite %0 with "transform.dialect"
}
sequence %arg0 {
^bb0(%arg1: !pdl.operation):
%1 = pdl_match @my_pattern in %arg1
// Use %1 as handle
}
}
Note that the pattern is expected to finish with a pdl.rewrite
terminator
that points to the custom rewriter named “transform.dialect”. The rewriter
actually does nothing, but the transform application will keep track of the
operations that matched the pattern.
This op is expected to contain pdl.pattern
operations and exactly one
another Transform dialect operation that gets executed with all patterns
available. This op is a possible top-level Transform IR op, the argument of
its entry block corresponds to either the root op of the payload IR or the
ops associated with its operand when provided.
Traits: NoTerminator, PossibleTopLevelTransformOpTrait, RecursiveSideEffects, SymbolTable
Interfaces: OpAsmOpInterface, TransformOpInterface
Operands: ¶
Operand | Description |
---|---|
root | PDL handle to an mlir::Operation * |
transform.yield
(::mlir::transform::YieldOp) ¶
Yields operation handles from a transform IR region
Syntax:
operation ::= `transform.yield` operands attr-dict (`:` type($operands)^)?
This terminator operation yields operation handles from regions of the transform IR ops back to the containing op. It is not itself associated with any transformation on the payload IR and is used for flow purposes only.
Traits: Terminator
Operands: ¶
Operand | Description |
---|---|
operands | any type |
TransformOpInterface (TransformOpInterface
) ¶
This interface is to be implemented by operations that identify transformations to be performed on other operations. The former are referred to as transform IR operations. The latter are referred to as payload IR operations. Such transform IR operations provide a fine-grain control mechanism over how transformations are applied by using and defining transform IR values, referred to as handles, that correspond to sets of operations in the payload IR. Transformations are applied starting from the operations identified by handles, but may affect other operations as well. Further restrictions may be imposed by flows that rely on transform IR operations to control transformations.
Methods: ¶
apply
¶
::mlir::LogicalResult apply(::mlir::transform::TransformResults &transformResults, ::mlir::transform::TransformState &state);
Applies the transformation represented by the current operation. This accepts as arguments the object that must be populated with results of the current transformation and a transformation state object that can be used for queries, e.g., to obtain the list of operations on which the transformation represented by the current op is targeted.
NOTE: This method must be implemented by the user.