This document describes a framework in MLIR in which to perform operation conversions between, and within dialects. This framework allows for transforming illegal operations to those supported by a provided conversion target, via a set of pattern-based operation rewriting patterns.
The dialect conversion framework consists of the following components:
When applying a conversion to a set of operations, there are several different conversion modes that may be selected from:
Partial Conversion
applyPartialConversion.Full Conversion
applyFullConversion.Analysis Conversion
applyAnalysisConversion.In all cases, the framework walks the operations in preorder, examining an op before the ops in any regions it has.
The conversion target is a formal definition of what is considered to be legal during the conversion process. The final operations generated by the conversion framework must be marked as legal on the ConversionTarget for the rewrite to be a success. Depending on the conversion mode, existing operations need not always be legal. Operations and dialects may be marked with any of the provided legality actions below:
Legal
Dynamic
arith.addi is only legal when operating on 32-bit integers.Illegal
An example conversion target is shown below:
struct MyTarget : public ConversionTarget { MyTarget(MLIRContext &ctx) : ConversionTarget(ctx) { //-------------------------------------------------------------------------- // Marking an operation as Legal: /// Mark all operations within the LLVM dialect are legal. addLegalDialect<LLVMDialect>(); /// Mark `arith.constant` op is always legal on this target. addLegalOp<arith::ConstantOp>(); //-------------------------------------------------------------------------- // Marking an operation as dynamically legal. /// Mark all operations within Affine dialect have dynamic legality /// constraints. addDynamicallyLegalDialect<AffineDialect>([](Operation *op) { ... }); /// Mark `std.return` as dynamically legal, but provide a specific legality /// callback. addDynamicallyLegalOp<ReturnOp>([](ReturnOp op) { ... }); /// Treat unknown operations, i.e. those without a legalization action /// directly set, as dynamically legal. markUnknownOpDynamicallyLegal([](Operation *op) { ... }); //-------------------------------------------------------------------------- // Marking an operation as illegal. /// All operations within the GPU dialect are illegal. addIllegalDialect<GPUDialect>(); /// Mark `std.br` and `std.cond_br` as illegal. addIllegalOp<BranchOp, CondBranchOp>(); } /// Implement the default legalization handler to handle operations marked as /// dynamically legal that were not provided with an explicit handler. bool isDynamicallyLegal(Operation *op) override { ... } };
In some cases, it may be desirable to mark entire regions as legal. This provides an additional granularity of context to the concept of “legal”. If an operation is marked recursively legal, either statically or dynamically, then all of the operations nested within are also considered legal even if they would otherwise be considered “illegal”. An operation can be marked via markOpRecursivelyLegal<>:
ConversionTarget &target = ...; /// The operation must first be marked as `Legal` or `Dynamic`. target.addLegalOp<MyOp>(...); target.addDynamicallyLegalOp<MySecondOp>(...); /// Mark the operation as always recursively legal. target.markOpRecursivelyLegal<MyOp>(); /// Mark optionally with a callback to allow selective marking. target.markOpRecursivelyLegal<MyOp, MySecondOp>([](Operation *op) { ... }); /// Mark optionally with a callback to allow selective marking. target.markOpRecursivelyLegal<MyOp>([](MyOp op) { ... });
After the conversion target has been defined, a set of legalization patterns must be provided to transform illegal operations into legal ones. The patterns supplied here have the same structure and restrictions as those described in the main Pattern documentation. The patterns provided do not need to generate operations that are directly legal on the target. The framework will automatically build a graph of conversions to convert non-legal operations into a set of legal ones.
As an example, say you define a target that supports one operation: foo.add. When providing the following patterns: [bar.add -> baz.add, baz.add -> foo.add], the framework will automatically detect that it can legalize bar.add -> foo.add even though a direct conversion does not exist. This means that you don’t have to define a direct legalization pattern for bar.add -> foo.add.
Along with the general RewritePattern classes, the conversion framework provides a special type of rewrite pattern that can be used when a pattern relies on interacting with constructs specific to the conversion process, the ConversionPattern. For example, the conversion process does not necessarily update operations in-place and instead creates a mapping of events such as replacements and erasures, and only applies them when the entire conversion process is successful. Certain classes of patterns rely on using the updated/remapped operands of an operation, such as when the types of results defined by an operation have changed. The general Rewrite Patterns can no longer be used in these situations, as the types of the operands of the operation being matched will not correspond with those expected by the user. This pattern provides, as an additional argument to the matchAndRewrite and rewrite methods, the list of operands that the operation should use after conversion. If an operand was the result of a non-converted operation, for example if it was already legal, the original operand is used. This means that the operands provided always have a 1-1 non-null correspondence with the operands on the operation. The original operands of the operation are still intact and may be inspected as normal. These patterns also utilize a special PatternRewriter, ConversionPatternRewriter, that provides special hooks for use with the conversion infrastructure.
struct MyConversionPattern : public ConversionPattern { /// The `matchAndRewrite` hooks on ConversionPatterns take an additional /// `operands` parameter, containing the remapped operands of the original /// operation. virtual LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const; };
The types of the remapped operands provided to a conversion pattern must be of a type expected by the pattern. The expected types of a pattern are determined by a provided TypeConverter. If no type converter is provided, the types of the remapped operands are expected to match the types of the original operands. If a type converter is provided, the types of the remapped operands are expected to be legal as determined by the converter. If the remapped operand types are not of an expected type, and a materialization to the expected type could not be performed, the pattern fails application before the matchAndRewrite hook is invoked. This ensures that patterns do not have to explicitly ensure type safety, or sanitize the types of the incoming remapped operands. More information on type conversion is detailed in the dedicated section below.
It is sometimes necessary as part of a conversion to convert the set types of being operated on. In these cases, a TypeConverter object may be defined that details how types should be converted when interfacing with a pattern. A TypeConverter may be used to convert the signatures of block arguments and regions, to define the expected inputs types of the pattern, and to reconcile type differences in general.
The TypeConverter contains several hooks for detailing how to convert types, and how to materialize conversions between types in various situations. The two main aspects of the TypeConverter are conversion and materialization.
A conversion describes how a given illegal source Type should be converted to N target types. If the source type is already “legal”, it should convert to itself. Type conversions are specified via the addConversion method described below.
A materialization describes how a set of values should be converted to a single value of a desired type. An important distinction with a conversion is that a materialization can produce IR, whereas a conversion cannot. These materializations are used by the conversion framework to ensure type safety during the conversion process. There are several types of materializations depending on the situation.
Argument Materialization
Source Materialization
Target Materialization
If a converted value is used by an operation that isn't converted, it needs a conversion back to the source type, hence source materialization; if an unconverted value is used by an operation that is being converted, it needs conversion to the target type, hence target materialization.
As noted above, the conversion process guarantees that the type contract of the IR is preserved during the conversion. This means that the types of value uses will not implicitly change during the conversion process. When the type of a value definition, either block argument or operation result, is being changed, the users of that definition must also be updated during the conversion process. If they aren't, a type conversion must be materialized to ensure that a value of the expected type is still present within the IR. If a target materialization is required, but cannot be performed, the pattern application fails. If a source materialization is required, but cannot be performed, the entire conversion process fails.
Several of the available hooks are detailed below:
class TypeConverter { public: /// Register a conversion function. A conversion function defines how a given /// source type should be converted. A conversion function must be convertible /// to any of the following forms(where `T` is a class derived from `Type`: /// * Optional<Type>(T) /// - This form represents a 1-1 type conversion. It should return nullptr /// or `llvm::None` to signify failure. If `llvm::None` is returned, the /// converter is allowed to try another conversion function to perform /// the conversion. /// * Optional<LogicalResult>(T, SmallVectorImpl<Type> &) /// - This form represents a 1-N type conversion. It should return /// `failure` or `llvm::None` to signify a failed conversion. If the new /// set of types is empty, the type is removed and any usages of the /// existing value are expected to be removed during conversion. If /// `llvm::None` is returned, the converter is allowed to try another /// conversion function to perform the conversion. /// * Optional<LogicalResult>(T, SmallVectorImpl<Type> &, ArrayRef<Type>) /// - This form represents a 1-N type conversion supporting recursive /// types. The first two arguments and the return value are the same as /// for the regular 1-N form. The third argument is contains is the /// "call stack" of the recursive conversion: it contains the list of /// types currently being converted, with the current type being the /// last one. If it is present more than once in the list, the /// conversion concerns a recursive type. /// Note: When attempting to convert a type, e.g. via 'convertType', the /// mostly recently added conversions will be invoked first. template <typename FnT, typename T = typename llvm::function_traits<FnT>::template arg_t<0>> void addConversion(FnT &&callback) { registerConversion(wrapCallback<T>(std::forward<FnT>(callback))); } /// Register a materialization function, which must be convertible to the /// following form: /// `Optional<Value> (OpBuilder &, T, ValueRange, Location)`, /// where `T` is any subclass of `Type`. /// This function is responsible for creating an operation, using the /// OpBuilder and Location provided, that "converts" a range of values into a /// single value of the given type `T`. It must return a Value of the /// converted type on success, an `llvm::None` if it failed but other /// materialization can be attempted, and `nullptr` on unrecoverable failure. /// It will only be called for (sub)types of `T`. /// /// This method registers a materialization that will be called when /// converting an illegal block argument type, to a legal type. template <typename FnT, typename T = typename llvm::function_traits<FnT>::template arg_t<1>> void addArgumentMaterialization(FnT &&callback) { argumentMaterializations.emplace_back( wrapMaterialization<T>(std::forward<FnT>(callback))); } /// This method registers a materialization that will be called when /// converting a legal type to an illegal source type. This is used when /// conversions to an illegal type must persist beyond the main conversion. template <typename FnT, typename T = typename llvm::function_traits<FnT>::template arg_t<1>> void addSourceMaterialization(FnT &&callback) { sourceMaterializations.emplace_back( wrapMaterialization<T>(std::forward<FnT>(callback))); } /// This method registers a materialization that will be called when /// converting type from an illegal, or source, type to a legal type. template <typename FnT, typename T = typename llvm::function_traits<FnT>::template arg_t<1>> void addTargetMaterialization(FnT &&callback) { targetMaterializations.emplace_back( wrapMaterialization<T>(std::forward<FnT>(callback))); } };
From the perspective of type conversion, the types of block arguments are a bit special. Throughout the conversion process, blocks may move between regions of different operations. Given this, the conversion of the types for blocks must be done explicitly via a conversion pattern. To convert the types of block arguments within a Region, a custom hook on the ConversionPatternRewriter must be invoked; convertRegionTypes. This hook uses a provided type converter to apply type conversions to all blocks within a given region, and all blocks that move into that region. As noted above, the conversions performed by this method use the argument materialization hook on the TypeConverter. This hook also takes an optional TypeConverter::SignatureConversion parameter that applies a custom conversion to the entry block of the region. The types of the entry block arguments are often tied semantically to details on the operation, e.g. FuncOp, AffineForOp, etc. To convert the signature of just the region entry block, and not any other blocks within the region, the applySignatureConversion hook may be used instead. A signature conversion, TypeConverter::SignatureConversion, can be built programmatically:
class SignatureConversion { public: /// Remap an input of the original signature with a new set of types. The /// new types are appended to the new signature conversion. void addInputs(unsigned origInputNo, ArrayRef<Type> types); /// Append new input types to the signature conversion, this should only be /// used if the new types are not intended to remap an existing input. void addInputs(ArrayRef<Type> types); /// Remap an input of the original signature with a range of types in the /// new signature. void remapInput(unsigned origInputNo, unsigned newInputNo, unsigned newInputCount = 1); /// Remap an input of the original signature to another `replacement` /// value. This drops the original argument. void remapInput(unsigned origInputNo, Value replacement); };
The TypeConverter provides several default utilities for signature conversion and legality checking: convertSignatureArgs/convertBlockSignature/isLegal(Region *|Type).
To debug the execution of the dialect conversion framework, -debug-only=dialect-conversion may be used. This command line flag activates LLVM's debug logging infrastructure solely for the conversion framework. The output is formatted as a tree structure, mirroring the structure of the conversion process. This output contains all of the actions performed by the rewriter, how generated operations get legalized, and why they fail.
Example output is shown below:
//===-------------------------------------------===//
Legalizing operation : 'std.return'(0x608000002e20) {
"std.return"() : () -> ()
* Fold {
} -> FAILURE : unable to fold
* Pattern : 'std.return -> ()' {
** Insert : 'spv.Return'(0x6070000453e0)
** Replace : 'std.return'(0x608000002e20)
//===-------------------------------------------===//
Legalizing operation : 'spv.Return'(0x6070000453e0) {
"spv.Return"() : () -> ()
} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//
} -> SUCCESS : pattern applied successfully
} -> SUCCESS
//===-------------------------------------------===//
This output is describing the legalization of an std.return operation. We first try to legalize by folding the operation, but that is unsuccessful for std.return. From there, a pattern is applied that replaces the std.return with a spv.Return. The newly generated spv.Return is then processed for legalization, but is found to already legal as per the target.