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.

A Transform IR value such as %0 may be associated with multiple payload operations. This is conceptually a set of operations and no assumptions should be made about the order of ops. Most Transform IR ops support operand values that are mapped to multiple operations. They usually apply the respective transformation for every mapped op (“batched execution”). Deviations from this convention are described in the documentation of Transform IR ops.

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 a Write effect.

  • 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 a Read 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 a Read.

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.

Execution Model ¶

The transformation starts at the specifed top-level transform IR operation and applies to some payload IR scope, identified by the payload IR op that contains the IR to transform. It is the responsibility of the user to properly select the scope and/or to avoid the transformations to modify the IR outside of the given scope. The top-level transform IR operation may contain further transform operations and execute them in the desired order.

Transformation application functions produce a tri-state status:

• success;
• recoverable (silencable) failure;
• irrecoverable failure.

Transformation container operations may intercept recoverable failures and perform the required recovery steps thus succeeding themselves. On the other hand, they must propagate irrecoverable failures. For such failures, the diagnostics are emitted immediately whereas their emission is postponed for recoverable failures. Transformation container operations may also fail to recover from a theoretically recoverable failure, in which case they are expected to emit the diagnostic and turn the failure into an irrecoverable one. A recoverable failure produced by applying the top-level transform IR operation is considered irrecoverable.

Transformation container operations are allowed to “step over” some nested operations if the application of some previous operation produced a failure. This can be conceptually thought of as having a global “recoverable error register” that is read/write accessed by each transform operation as a side effect. The transformation is skipped if the register already contains an error description, and the control flow proceeds to the following operation.

Handle Invalidation ¶

The execution model of the transform dialect expects that a payload IR operation is associated with at most one transform IR handle. This avoids the situation when a handle to an operation outlives the operation itself that can be erased during a transformation triggered through another handle.

Handles pointing to operations nested in each other are allowed to co-exist in the transform IR. However, a transform IR operation that consumes such a handle automatically invalidates all the other handles that are associated with operations nested in the operations associated with the consumed handle. Any use of the invalidated handle results in undefined behavior since the payload IR operations associated with it are likely to have been mutated or erased. The mere fact of the handle being invalidated does not trigger undefined behavior, only its appearance as an operand does. Invalidation applies to the entire handle, even if some of the payload IR operations associated with it are not nested in payload IR operations associated with another, consumed handle.

Note: the restriction on two handles not pointing to the same operation may be relaxed in the future to follow the invalidation model for nested operation.

The Transform dialect infrastructure has the capability of checking whether the transform IR op operand is invalidated before applying the transformation. However, such a check is computationally expensive and must be enabled explicitly through TransformOptions. Additionally, the transform-dialect-check-uses pass emits warnings when a handle may be used after it has been consumed, but does so abstractly, without processing the payload IR.

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 transformation 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 transform 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.alternatives (::mlir::transform::AlternativesOp) ¶

Attempts sequences of transforms until one succeeds

Syntax:

operation ::= `transform.alternatives` ($scope^)? (`->` type($results)^)? attr-dict-with-keyword regions

This op may have an arbitrary number of regions, each of which represents a sequence of transform operations to be applied to the same payload IR. The regions are visited in order of appearance, and transforms in them are applied in their respective order of appearance. If one of these transforms fails to apply, the remaining ops in the same region are skipped an the next region is attempted. If all transformations in a region succeed, the remaining regions are skipped and the entire “alternatives” transformation succeeds. If all regions contained a failing transformation, the entire “alternatives” transformation fails.

It is up to the nested operations to define which errors are “recoverable” (or “silenceable”) and allow another alternatives to be attempted, and which errors should be propagated without attempting the other alternatives.

The single operand of this operation is the scope in which the alternative transformation sequences are attempted, that is, an operation in the payload IR that contains all the other operations that may be modified by the transformations. The scope operation must be isolated from above. There is no check that the transforms are indeed scoped as their “apply” methods can be arbitrarily complex. Therefore it is the responsibility of the user to ensure that the transforms are scoped correctly, or to produce an irrecoverable error and thus abort the execution without attempting the remaining alternatives. Note that the payload IR outside of the given scope is not necessarily in the valid state, or even accessible to the transformation.

The changes to the IR within the scope performed by transforms in the failed alternative region are reverted before attempting the next region. Practically, this is achieved by cloning the scope. Therefore it is advised to limit the scope as much as possible and place the most likely alternatives early in the region list. The operation is also isolated from above and requires rediscovering the operations within the given scope to avoid additional handle invalidation. The latter restriction may be lifted in the future.

Each of the regions may yield transform IR handles. The handles of the first successful alternative region are returned as the results of the “alternatives” op. Therefore, each alternative region must yield the same number of results, which should also match the number and the types of the “alternatives” op results.

Remark: this op allows one to implement a simple “try” construct as follows:

%result = transform.alternatives %scope {
^bb0(%arg0: !pdl.operation):
  // Try a fallible transformation.
  %0 = transform.fallible %arg0 // ...
  // If succeeded, yield the the result of the transformation.
  transform.yield %0 : !pdl.operation
}, {
^bb0(%arg0: !pdl.operation):
  // Otherwise, the second alternative is tried and it always succeeds by
  // returning the original handle.
  transform.yield %arg0 : !pdl.operation
}

Traits: FunctionalStyleTransformOpTrait, IsolatedFromAbove, PossibleTopLevelTransformOpTrait, SingleBlockImplicitTerminator<::mlir::transform::YieldOp>

Interfaces: MemoryEffectsOpInterface, RegionBranchOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.foreach (::mlir::transform::ForeachOp) ¶

Executes the body for each payload op

Syntax:

operation ::= `transform.foreach` $target (`->` type($results)^)? $body attr-dict

This op has exactly one region with exactly one block (“body”). The body is executed for each payload op that is associated to the target operand in an unbatched fashion. I.e., the block argument (“iteration variable”) is always mapped to exactly one payload op.

This op always reads the target handle. Furthermore, it consumes the handle if there is a transform op in the body that consumes the iteration variable. This op does not return anything.

The transformations inside the body are applied in order of their appearance. 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.

This op generates as many handles as the terminating YieldOp has operands. For each result, the payload ops of the corresponding YieldOp operand are merged and mapped to the same resulting handle.

Traits: SingleBlockImplicitTerminator<::mlir::transform::YieldOp>

Interfaces: MemoryEffectOpInterface, RegionBranchOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.get_closest_isolated_parent (::mlir::transform::GetClosestIsolatedParentOp) ¶

Gets handles to the closest isolated-from-above parents

Syntax:

operation ::= `transform.get_closest_isolated_parent` $target attr-dict

The handles defined by this Transform op correspond to the closest isolated from above ancestor of the Payload IR operations associated with its operand. If any of the given Payload IR ops has no such parent (unlikely as there usually is a top-level ModuleOp), the transformation is considered to have failed.

Ancestor ops follow the same order as the ops associated with the operand, except for potential duplicates (multiple Payload IR ops associated with the operand have the same parent) for which the ancestor will only be listed once for the first time it occurs. For example, given the list “(childof(A), childof(B), childof(B), childof(A), childof(B))”, the resulting list will be just “(A, B)”. Note that no other semantic ordering is applied, e.g., “B” may itself be a parent of “A”. This may have an impact on the further transformation applied to the handle produced here.

Traits: NavigationTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.merge_handles (::mlir::transform::MergeHandlesOp) ¶

Merges handles into one pointing to the union of payload ops

Syntax:

operation ::= `transform.merge_handles` ($deduplicate^)? $handles attr-dict

Creates a new Transform IR handle value that points to the same Payload IR operations as the operand handles. The Payload IR operations are listed in the same order as they are in the operand handles, grouped by operand handle, e.g., all Payload IR operations associated with the first handle come first, then all Payload IR operations associated with the second handle and so on. If deduplicate is set, do not add the given Payload IR operation more than once to the final list regardless of it coming from the same or different handles. Consumes the operands and produces a new handle.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

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 associated 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 transformation 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: ¶

Operands: ¶

Results: ¶

transform.replicate (::mlir::transform::ReplicateOp) ¶

Lists payload ops multiple times in the new handle

Syntax:

operation ::= `transform.replicate` `num` `(` $pattern `)` $handles custom<PDLOpTypedResults>(type($replicated), ref($handles)) attr-dict

Produces a new handle associated with a list of payload IR ops that is computed by repeating the list of payload IR ops associated with the operand handle as many times as the “pattern” handle has associated operations. For example, if pattern is associated with [op1, op2] and the operand handle is associated with [op3, op4, op5], the resulting handle will be associated with [op3, op4, op5, op3, op4, op5].

This transformation is useful to “align” the sizes of payload IR lists before a transformation that expects, e.g., identically-sized lists. For example, a transformation may be parameterized by same notional per-target size computed at runtime and supplied as another handle, the replication allows this size to be computed only once and used for every target instead of replicating the computation itself.

Note that it is undesirable to pass a handle with duplicate operations to an operation that consumes the handle. Handle consumption often indicates that the associated payload IR ops are destroyed, so having the same op listed more than once will lead to double-free. Single-operand MergeHandlesOp may be used to deduplicate the associated list of payload IR ops when necessary. Furthermore, a combination of ReplicateOp and MergeHandlesOp can be used to construct arbitrary lists with repetitions.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

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, RegionBranchOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

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: ¶

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: ¶

Bufferization Transform Operations ¶

transform.bufferization.one_shot_bufferize (::mlir::transform::OneShotBufferizeOp) ¶

Syntax:

operation ::= `transform.bufferization.one_shot_bufferize` $target attr-dict

Indicates that the given target op should be bufferized with One-Shot Bufferize. The bufferization can be configured with various attributes that corresponding to options in BufferizationOptions and the one-shot-bufferize pass. More information can be found in the pass documentation.

If target_is_module is set, target must be a module. In that case the target handle can be reused by other transform ops. When bufferizing other ops, the target handled is freed after bufferization and can no longer be used.

Note: Only ops that implement BufferizableOpInterface are bufferized. All other ops are ignored if allow_unknown_ops. If allow_unknown_ops is unset, this transform fails when an unknown/non-bufferizable op is found. Many ops implement BufferizableOpInterface via an external model. These external models must be registered when applying this transform op; otherwise, said ops would be considered non-bufferizable.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Loop (SCF) Transform Operations ¶

transform.loop.get_parent_for (::mlir::transform::GetParentForOp) ¶

Gets a handle to the parent ‘for’ loop of the given operation

Syntax:

operation ::= `transform.loop.get_parent_for` $target attr-dict

Produces a handle to the n-th (default 1) parent scf.for loop for each Payload IR operation associated with the operand. Fails if such a loop cannot be found. The list of operations associated with the handle contains parent operations in the same order as the list associated with the operand, except for operations that are parents to more than one input which are only present once.

Traits: NavigationTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.loop.outline (::mlir::transform::LoopOutlineOp) ¶

Outlines a loop into a named function

Syntax:

operation ::= `transform.loop.outline` $target attr-dict

Moves the loop into a separate function with the specified name and replaces the loop in the Payload IR with a call to that function. Takes care of forwarding values that are used in the loop as function arguments. If the operand is associated with more than one loop, each loop will be outlined into a separate function. The provided name is used as a base for forming actual function names following SymbolTable auto-renaming scheme to avoid duplicate symbols. Expects that all ops in the Payload IR have a SymbolTable ancestor (typically true because of the top-level module). Returns the handle to the list of outlined functions in the same order as the operand handle.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.loop.peel (::mlir::transform::LoopPeelOp) ¶

Peels the last iteration of the loop

Syntax:

operation ::= `transform.loop.peel` $target attr-dict

Updates the given loop so that its step evenly divides its range and puts the remaining iteration into a separate loop or a conditional.

In the absence of sufficient static information, this op may peel a loop, even if the step always divides the range evenly at runtime.

Return modes ¶

This operation ignores non-scf::ForOp ops and drops them in the return.

This operation always succeeds and returns the scf::ForOp with the postcondition: “the loop trip count is divisible by the step”. This operation may return the same unmodified loop handle when peeling did not modify the IR (i.e. the loop trip count was already divisible).

Note that even though the Payload IR modification may be performed in-place, this operation consumes the operand handle and produces a new one.

TODO: Return both the peeled loop and the remainder loop.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.loop.pipeline (::mlir::transform::LoopPipelineOp) ¶

Applies software pipelining to the loop

Syntax:

operation ::= `transform.loop.pipeline` $target attr-dict

Transforms the given loops one by one to achieve software pipelining for each of them. That is, performs some amount of reads from memory before the loop rather than inside the loop, the same amount of writes into memory after the loop, and updates each iteration to read the data for a following iteration rather than the current one.

The amount is specified by the attributes.

The values read and about to be stored are transferred as loop iteration arguments. Currently supports memref and vector transfer operations as memory reads/writes.

Return modes ¶

This operation ignores non-scf::For ops and drops them in the return. If all the operations referred to by the target PDLOperation pipeline properly, the transform succeeds. Otherwise the transform silently fails. The return handle points to only the subset of successfully produced pipelined loops, which can be empty.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.loop.unroll (::mlir::transform::LoopUnrollOp) ¶

Unrolls the given loop with the given unroll factor

Syntax:

operation ::= `transform.loop.unroll` $target attr-dict

Unrolls each loop associated with the given handle to have up to the given number of loop body copies per iteration. If the unroll factor is larger than the loop trip count, the latter is used as the unroll factor instead.

Return modes ¶

This operation ignores non-scf::For ops and drops them in the return. If all the operations referred to by the target PDLOperation unroll properly, the transform succeeds. Otherwise the transform silently fails.

Does not return handles as the operation may result in the loop being removed after a full unrolling.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Structured (Linalg) Transform Operations ¶

transform.structured.decompose (::mlir::transform::DecomposeOp) ¶

Syntax:

operation ::= `transform.structured.decompose` $target attr-dict

Decomposes named complex operations, such as higher-dimensional (depthwise) convolutions, into combinations of lower-dimensional equivalents when possible.

Return modes ¶

This operation ignores non-Linalg ops and drops them in the return. If all the operations referred to by the target PDLOperation decompose properly, the transform succeeds. Otherwise the transform silently fails. The return handle points to only the subset of successfully produced computational operations, which can be empty.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.structured.fuse_into_containing_op (::mlir::transform::FuseIntoContainingOp) ¶

Fuses the `producer_op` into the `containing_op`. Only producers with a
single result are supported at the moment. Returns a handle to the fused
ops.

The producer is typically a slice of a tileable op (i.e., implements
TilingInterface). In that case, this transform computes the accessed
producer slice inside of the containing op ("tile and fuse"). Otherwise,
the entire producer is cloned inside the containing op ("clone and fuse").

The containing op handle must be associated with exactly one payload op. The
producer op handle may be associated with multiple payload ops. This
transform fuses producers one-by-one, always picking an unspecified producer
that has at least one use inside the containing op among the 
producers.

Note: If a producer has multiple uses inside the containing op, it is
currently tiled and/or cloned multiple times into the containing op.
TODO: Reuse already fused OpResults instead of tiling/cloning a second time
when possible. Fuse producers according to a topological sorting to achieve
the largest amount of reuse.

#### Return modes

If at least one producer could not be fused, this operation fails silently.
This is the case when tiling fails or when no producer op could be found
among the remaining producers that has at least one use within the
containing op. I.e., "producers" that are not consumed within the containing
op are rejected by this operation. This operation reads and frees the
producer handle. It reads the containing op handle.

Syntax:

operation ::= `transform.structured.fuse_into_containing_op` $producer_op `into` $containing_op attr-dict

Fuse a producer into a containing operation. Interfaces: TransformOpInterface

Operands: ¶

Results: ¶

transform.structured.fuse (::mlir::transform::FuseOp) ¶

Tiles the operations pointed to by the target handle and fuses their producers greedily using the options provided as attributes.

Traits: FunctionalStyleTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.generalize (::mlir::transform::GeneralizeOp) ¶

Syntax:

operation ::= `transform.structured.generalize` $target attr-dict

Transforms a named structured operation into the generic form with the explicit attached region.

Return modes ¶

This operation ignores non-Linalg ops and drops them in the return. If all the operations referred to by the target PDLOperation generalize properly, the transform succeeds. Otherwise the transform silently fails. The return handle points to only the subset of successfully produced equivalent generic operations, which can be empty or contain the original ops if they were already in generic form.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.structured.interchange (::mlir::transform::InterchangeOp) ¶

Syntax:

operation ::= `transform.structured.interchange` $target attr-dict

Interchanges the iterators of the operations pointed to by the target handle using the iterator interchange attribute.

Return modes ¶

This operation ignores non-linalg::Generic ops and drops them in the return. This operation fails if the interchange attribute is invalid. If all the operations referred to by the target PDLOperation interchange properly, the transform succeeds. If any interchange fails, the transform definitely fails. The return handle points to only the subset of successfully produced interchanged operations, which can be empty.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.match (::mlir::transform::MatchOp) ¶

Syntax:

operation ::= `transform.structured.match` (`ops` `{` $ops^ `}`)?
              (`interface` `{` $interface^ `}`)?
              (`attribute` `{` $attribute^ `}`)?
              `in` $target attr-dict

Match op with the specified constraints, within the target op.

The following constraints are supported:

• interface: an optional MatchInterfaceEnum specifying an enum representation for an interface to target.
• ops: an optional StrArrayAttr specifying the concrete name of an op. Multiple names can be specified. Matched ops must have one of specified names.
• attribute: an optional Str specifying the name of an attribute that matched ops must have.

Note: Only ops that satisfy all specified constraints are matched.

TODO: Extend with regions to allow a limited form of constraints.

Return modes ¶

This op traverses the ops nested under target and returns the handles to all the operations that match the requirements.

This op fails if the target is not a handle to exactly one operation. Otherwise it succeeds.

This operation does not consume the target handle and produces new handles: it is a navigation op.

Traits: NavigationTransformOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.multitile_sizes (::mlir::transform::MultiTileSizesOp) ¶

Syntax:

operation ::= `transform.structured.multitile_sizes` $target attr-dict

Emits the IR computing the tile sizes s1 and s2 such that:

• there exists a combination of n tiles of size s1 and m tiles of size s2 that covers the entirety of the iteration space dimension of the target structured op;
• s1, s2 is less than or equal to target_size;
• s1 and s2 are divisible by `divisor.

For example, for a dimension of size 54 with target size 12 and divisor 2, this can emit the IR computing the tile size 10, used for 3 tiles, and 12, used for 2 tiles, totally 103 + 122 = 54. Note that when the divisor does not divide the original dimension size, it is impossible to compute such tile sizes. An assertion is emitted to guard against this in the dynamic case.

Expects the target size and the divisor to be strictly positive. Folds the IR as much as possible, normally obtaining constant sizes and numbers of tiles for a statically known dimension.

This does not consume the target handle and produces three handles each pointing to single-result index-typed operations (which may be arithmetic constant operations) defining the two respective tile sizes and the product of the first tile size with the number of tiles of that size (useful for splitting the iteration space).

This operation composes with the regular tiling when applied per-dimension:

%sz1, %sz2, %split = structured.multitile_sizes %target
                     { target_size = 10, dimension = 1 }
%low, %high = structured.split %target after %split { dimension = 1 }
%tiled_low = structured.tile %low [0, %sz1]
%tiled_high = structured.tile %high [0, %sz2]
%common = merge_handles %tiled_low, %tiled_high

%sz3, %sz4, %split = structured.multitile_size %target
                     { target_size = 42, dimension = 0 }
%sz3r, %sz4r, %splitr = replicate num(%common) %sz3, %sz4, %splitr
structured.split %common after %splitr { dimension = 0 }
// ...

Traits: TransformEachOpTrait

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.pad (::mlir::transform::PadOp) ¶

Syntax:

operation ::= `transform.structured.pad` $target attr-dict

Pads the operations pointed to by the target handle using the options provides as operation attributes.

Return modes ¶

This operation ignores non-Linalg ops and drops them in the return. This operation may produce a definiteFailure if the padding fails for any reason. If all the operations referred to by the target PDLOperation pad properly, the transform succeeds. Otherwise the transform silently fails. The return handle points to only the subset of successfully produced padded operations, which can be empty.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.promote (::mlir::transform::PromoteOp) ¶

Syntax:

operation ::= `transform.structured.promote` $target attr-dict

Promotes the specified operands of the target into a separate memory buffer.

At this point, this transform does not allow customizing alloc/dealloc functions nor the behavior on copy in/out operations.

Return modes ¶

This operation applies to a single Linalg op that satisfies the promoteSubviewsPrecondition, otherwise it fails.

If the operations referred to by the target PDLOperation promote properly, the transform succeeds.

When successful, the return handle points to the $target operation that was modified inplace.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.scalarize (::mlir::transform::ScalarizeOp) ¶

Syntax:

operation ::= `transform.structured.scalarize` $target attr-dict

Indicates that ops of a specific kind in the given function should be scalarized (i.e. their dynamic dimensions tiled by 1).

Return modes: ¶

This operation ignores non-Linalg ops and drops them in the return. This operation produces definiteFailure if the scalarization fails for any reason. If all the operations referred to by the target PDLOperation scalarize properly, the transform succeeds. Otherwise the transform silently fails.

The return handle points to only the subset of successfully produced tiled-by-1 operations, which can be empty.

This operation does not return handles to the tiled loop. We make this design choice because it is hard to know ahead of time the number of loops that will be produced (it depends on the number of dynamic dimensions after multiple transformations have been applied). Loops can always be recovered by navigating from the tiled operations if needed.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Operands: ¶

Results: ¶

transform.structured.split (::mlir::transform::SplitOp) ¶

Indicates that the given target op should be split into two complementary parts, which combined cover the entire iteration domain of the original op. The split is performed along the iteration space dimension provided as attribute. In case of dimension overflow, the transformation fails. The split is performed at the dimension iterator value specified as either the static split point attribute when it is known at transform IR construction time or as the handle to an operation producing a single index-typed value when it is computed by payload IR. In the latter case, the static split point must be set to ShapedType::kDynamicSize and the dynamic size handle must point to as many value-producing operations as there are structured operations pointed to by the target handle.

The operation consumes the target handle, but preserves the split point handle if provided. It produces two new handles pointing to the two parts of the structured op after splitting, in the same order as the target operand, with the first handle corresponding to the part with lower iteration space indices.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.split_reduction (::mlir::transform::SplitReductionOp) ¶

Syntax:

operation ::= `transform.structured.split_reduction` $target attr-dict

Indicates that the given target op should be transformed with the splitReduction transformation and split factor provided as attribute.

The splitReduction transformation splits the first single linalg op reduction into a parallel and reduction dimension. A new linalg.generic op is created to perform the rest of the reduction.

The transformation supports different configurations attributes:

• split_factor: the factor by which to split (i.e. the size of the remaining reduction after splitting).
• insert_split_dimension: the dimension in the temporary tensor into which the new parallel dimension is inserted.
• use_scaling_algorithm: whether to use a scaling based formulation that does not create an ExpandShapeOp (default: do not use scaling)
• use_alloc: whether to use an alloc op to allocate the temporary tensor (default: do not use alloc op)

Return modes ¶

This operation ignores non-Linalg ops and drops them in the return. This operation produces definiteFailure if the splitting fails for any reason.

If all the operations referred to by the target PDLOperation split properly, the transform succeeds. Otherwise the transform silently fails. The 4 returned handles points to only the subset of successfully produced computational operations, which can all be empty. This 4 returned handles point to:

• the init op (or tensor_alloc op if use_alloc = true),
• the fill op used to initialize the neutral element,
• the split op and
• the result-combining op.

Example (default: use_scaling_algorithm = false, use_alloc = false): ¶

  %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
                                        affine_map<(d0) -> ()>],
        iterator_types = ["reduction"]}
  ins(%in : tensor<32xf32>)
  outs(%out : tensor<f32>) {
  ^bb0(%arg1: f32, %arg2: f32):
    %y = arith.addf %arg1, %arg2 : f32
    linalg.yield %y : f32
  } -> tensor<f32>

is split into:

  %cst = arith.constant 0.000000e+00 : f32
  %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
  %1 = linalg.init_tensor [4] : tensor<4xf32>
  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32>
  %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                                        affine_map<(d0, d1) -> (d0)>],
    iterator_types = ["parallel", "reduction"]}
    ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
    ^bb0(%arg3: f32, %arg5: f32):
    %5 = arith.addf %arg3, %arg4 : f32
    linalg.yield %5 : f32
  } -> tensor<4xf32>
  %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
                                        affine_map<(d0) -> ()>],
    iterator_types = ["reduction"]}
    ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
    ^bb0(%arg3: f32, %arg4: f32):
    %5 = arith.addf %arg3, %arg4 : f32
    linalg.yield %5 : f32
  } -> tensor<f32>

Example (use_scaling_algorithm = true, use_alloc = true): ¶

Instead of introducing an ExpandShapeOp, this scaling-based implementation rewrites a reduction dimension k into k * split_factor + kk. The dimension kk is added as an extra parallel dimension to the intermediate output tensor at position insert_split_dimension.

Consider a minimal example where k is reduced: O(i, j) += I(i, j, k) Assume i=3, j=5, k=128, split_factor=16 and insert_split_dimension=0. The compute is rewritten as: a. O_i(kk, i, j) += I(i, j, 16 * k + kk) b. O(i, j) += O_i(kk, i, j) The intermediate tensor O_i is of shape (128/16)x3x5 == 8x3x5.

Example: ¶

 %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
   outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>

Is transformed to:

 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
 #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
 #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
 #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
 %0 = linalg.init_tensor [16, 32, 64] : tensor<16x32x64xf32>
 %cst = arith.constant 0.000000e+00 : f32
 %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
    tensor<16x32x64xf32>
 %2 = linalg.init_tensor [64, 4] : tensor<64x4xi1>

 %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
   iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
   ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>)
   outs(%1 : tensor<16x32x64xf32>) {
     ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
       %5 = arith.mulf %arg3, %arg4 : f32
       %6 = arith.addf %arg6, %5 : f32
       linalg.yield %6 : f32
 } -> tensor<16x32x64xf32>

 %4 = linalg.generic {indexing_maps = [#map4, #map5],
   iterator_types = ["parallel", "parallel", "reduction"]}
   ins(%3 : tensor<16x32x64xf32>)
   outs(%C : tensor<16x32xf32>) {
     ^bb0(%arg3: f32, %arg4: f32):
       %5 = arith.addf %arg3, %arg4 : f32
       linalg.yield %5 : f32
 } -> tensor<16x32xf32>

 return %4 : tensor<16x32xf32>

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.tile (::mlir::transform::TileOp) ¶

Indicates that the given target op should be tiled with the given sizes. This transform generates a loop nest with a smaller (“tiled”) target operation in its body. Currently limited to LinalgOps.

Tile sizes may be known at transformation time, in which case they are expected to be provided in the static_size attribute, or not, in which case the tile value must be computed by the payload IR and the handle to the operation computing it must be provided through dynamic_sizes. When the sizes are not known statically, the corresponding entry in the static_sizes attribute must be set to ShapedType::kDynamicSize. Only the dynamic sizes must be provided in dynamic_sizes, i.e., there should be as many handles as ShapedType::kDynamicSize values in the static_sizes attribute. A static size of 0 indicates that the dimension should not be tiled. No loop will be generated for such dimensions. If all tile sizes are 0, this transform is effectively a no-op.

This op returns handles to the tiled op (in the generated loop nest) and the generated loops. The number of loops is the number of tile sizes that are statically known to be non-zero.

Return modes ¶

On success, the resulting handles are associated with co-indexed lists of tiled operations and loops around them.

This operation only supports Linalg ops and produces a silenceable failure if the input contains any non-Linalg ops. The ops preceding it in the list associated with the target handle will have been tiled.

This operation produces a silenceable failure if the dynamic_sizes handles are associated with lists of payload operations of a size different than that of the list associated with the target handle.

If the internal implementation of tiling for any of the operations fails, produces a definite failure.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.tile_to_foreach_thread_op (::mlir::transform::TileToForeachThreadOp) ¶

Syntax:

operation ::= `transform.structured.tile_to_foreach_thread_op` $target  (`num_threads` $num_threads^) : (`tile_sizes` $tile_sizes)?
              (`(` `mapped` `to` `dims` $thread_dim_mapping^ `)`)? attr-dict

Tile a TilingInterface op to a tiled scf.foreach_thread. Tiling is applied by either specifying num_threads or tile_size. If num_threads is specified, then the tile size for each dimension i is calculated dynamically via ceilDiv(dimSize[i], num_threads[i]). If non-empty, the thread_dim_mapping is added as an attribute to the resulting scf.foreach_thread. Zero tile sizes indicate that the dimension is not tiled and can be thought of as tiling by the full size of data. It is the user’s responsibility to ensure that num_threads/tile_sizes is a valid tiling specification (i.e. that only tiles parallel dimensions, e.g. in the Linalg case).

Return modes ¶

This operation ignores ops that do not implement the TilingInterface and drops them in the return.

If all the operations referred to by the target PDLOperation tile successfully, the transform succeeds. Otherwise the transform silently fails.

The two returned handles point to only the subset of successfully produced tiled operations, which can all be empty.

These two returned handles point to:

• the new scf.foreach_thread op,
• the tiled op that implements TilingInterface.

Example using num_threads ¶

%0 = pdl_match @match_matmul in %arg1    
%3:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20]

Example using tile_sizes ¶

%0 = pdl_match @match_matmul in %arg1    
%3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 20, 0]

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

transform.structured.vectorize (::mlir::transform::VectorizeOp) ¶

Syntax:

operation ::= `transform.structured.vectorize` $target attr-dict

Indicates that the given target op all the ops it contains should be vectorized with the configuration specified by the attributes of this op. This vectorization only handles structured ops that operate on shaped types and does not vectorize loops or straight-line. Internally, it applies a set of rewrite patterns, some of which enable vectorization and some of which clean up the results. Therefore, it can only be applied to an op with the “isolated from above property”. If finer granularity is required, it can be achieved by outlining the target part of the payload IR into, e.g., a function, performing the transformation, and inlining it back. This transformation only fails if the entire pattern rewriting failed, i.e., it does not fail when no ops were vectorized.

Note that this transformation is invalidating the handles to any payload IR operation that is contained inside the vectorization target.

Return modes: ¶

This operation produces definiteFailure if vectorization fails for any reason. The operation always returns the handle to the target op that is expected to be isolated from above.

Traits: FunctionalStyleTransformOpTrait, TransformEachOpTrait

Interfaces: MemoryEffectsOpInterface, TransformOpInterface

Attributes: ¶

Operands: ¶

Results: ¶

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::DiagnosedSilenceableFailure 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. Returns a special status object indicating whether the transformation succeeded or failed, and, if it failed, whether the failure is recoverable or not.

NOTE: This method must be implemented by the user.