llvm-project/mlir/docs/DialectConversion.md

# Dialect Conversion

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:

*   A [Conversion Target](#conversion-target)
*   A set of [Rewrite Patterns](#rewrite-pattern-specification)
*   A [Type Converter](#type-conversion) (Optional)

[TOC]

## Modes of Conversion

When applying a conversion to a set of operations, there are several different
conversion modes that may be selected from:

*   Partial Conversion

    -   A partial conversion will legalize as many operations to the target as
        possible, but will allow pre-existing operations that were not
        explicitly marked as "illegal" to remain unconverted. This allows for
        partially lowering parts of the input in the presence of unknown
        operations.
    -   A partial conversion can be applied via `applyPartialConversion`.

*   Full Conversion

    -   A full conversion legalizes all input operations, and is only successful
        if all operations are properly legalized to the given conversion target.
        This ensures that only known operations will exist after the conversion
        process.
    -   A full conversion can be applied via `applyFullConversion`.

*   Analysis Conversion

    -   An analysis conversion will analyze which operations are legalizable to
        the given conversion target if a conversion were to be applied. This is
        done by performing a 'partial' conversion and recording which operations
        would have been successfully converted if successful. Note that no
        rewrites, or transformations, are actually applied to the input
        operations.
    -   An analysis conversion can be applied via `applyAnalysisConversion`.

In all cases, the framework walks the operations in preorder, examining an op
before the ops in any regions it has.

## Conversion Target

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

    -   This action signals that every instance of a given operation is legal,
        i.e. any combination of attributes, operands, types, etc. are valid.

*   Dynamic

    -   This action signals that only some instances of a given operation are
        legal. This allows for defining fine-tune constraints, e.g. saying that
        `arith.addi` is only legal when operating on 32-bit integers.

*   Illegal

    -   This action signals that no instance of a given operation is legal.
        Operations marked as "illegal" must always be converted for the
        conversion to be successful. This action also allows for selectively
        marking specific operations as illegal in an otherwise legal dialect.

Operations and dialects that are neither explicitly marked legal nor illegal are
separate from the above ("unknown" operations) and are treated differently, for
example, for the purposes of partial conversion as mentioned above.

An example conversion target is shown below:

```c++
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<affine::AffineDialect>(
        [](Operation *op) { ... });

    /// Mark `func.return` as dynamically legal, but provide a specific legality
    /// callback.
    addDynamicallyLegalOp<func::ReturnOp>([](func::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 `cf.br` and `cf.cond_br` as illegal.
    addIllegalOp<cf::BranchOp, cf::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 { ... }
};
```

### Recursive Legality

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

```c++
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) { ... });
```

## Rewrite Pattern Specification

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 similar restrictions as those
described in the main [Pattern](PatternRewriter.md) 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`.

### Conversion Patterns

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`.

#### Remapped Operands / Adaptor
Conversion patterns have an additional `operands` / `adaptor` argument for the
`matchAndRewrite` method. These operands correspond to the most recent
replacement values of the respective operands of the matched operation.

```c++
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;
};
```

Example:

```mlir
%0 = "test.foo"() : () -> i1  // matched by pattern A
"test.bar"(%0) : (i1) -> ()   // matched by pattern B
```

Let's assume that the two patterns are applied back-to-back: first, pattern A
replaces `"test.foo"` with `"test.qux"`, an op that has a different result
type. The dialect conversion infrastructure has special support for such
type-changing IR modifications.

```mlir
%0 = "test.qux"() : () -> i2
%r = builtin.unrealized_conversion_cast %0 : i2 to i1
"test.bar"(%r) : (i1) -> ()
```

Simply swapping out the operand of `"test.bar"` during the `replaceOp` call
would be unsafe, because that would change the type of operand and, therefore,
potentially the semantics of the operation. Instead, the dialect conversion
driver (conceptually) inserts a `builtin.unrealized_conversion_cast` op that
connects the newly created `"test.qux"` op with the `"test.bar"` op, without
changing the types of the latter one.

Now, the second pattern B is applied. The `operands` argument contains an SSA
value with the most recent replacement type (`%0` with type `i2`), whereas
querying the operand from the matched op still returns an SSA value with the
original operand type `i1`.

Note: If the conversion pattern is instantiated with a type converter, the
`operands` argument contains SSA values whose types match the legalized operand
types as per the type converter. See [Type Safety](#type-safety) for more
details.

Note: The dialect conversion framework does not guarantee the presence of any
particular value in the `operands` argument. The only thing that's guaranteed
is the type of the `operands` SSA values. E.g., instead of the actual
replacement values supplied to a `replaceOp` API call, `operands` may contain
results of transitory `builtin.unrealized_conversion_cast` ops that were
inserted by the conversion driver but typically fold away again throughout the
conversion process.

#### Immediate vs. Delayed IR Modification

The dialect conversion driver can operate in two modes: (a) rollback mode
(default) and (b) no-rollback mode. This can be controlled by
`ConversionConfig::allowPatternRollback`. When running in rollback mode, the
driver is able backtrack and roll back already applied patterns when the
current legalization path (sequence of pattern applications) gets stuck with
unlegalizable operations.

When running in no-rollback mode, all IR modifications such as op replacement,
op erasure, op insertion or in-place op modification are applied immediately.

When running in rollback mode, certain IR modifications are delayed to the end
of the conversion process. For example, a `ConversionPatternRewriter::eraseOp`
API call does not immediately erase the op, but just marks it for erasure. The
op will stay visible to patterns and IR traversals throughout the conversion
process. As another example, `replaceOp` and `replaceAllUsesWith` does not
immediately update users of the original SSA values. This step is also delayed
to the end of the conversion process.

Delaying certain IR modifications has two benefits: (1) pattern rollback is
simpler because fewer IR modifications must be rolled back, (2) pointers of
erased operations / blocks are preserved upon rollback, and (3) patterns can
still access/traverse the original IR to some degree. However, additional
bookkeeping in the form of complex internal C++ data structures is required to
support pattern rollback. Running in rollback mode has a significant toll on
compilation time, is error-prone and makes debugging conversion passes more
complicated. Therefore, programmers are encouraged to run in no-rollback mode
when possible.

The following table gives an overview of which IR changes are applied in a
delayed fasion in rollback mode.

| Type                                                    | Rollback Mode     | No-rollback Mode |
| ------------------------------------------------------- | ----------------- | ---------------- |
| Op Insertion / Movement (`create`/`insert`)             | Immediate         | Immediate        |
| Op Replacement (`replaceOp`)                            | Delayed           | Immediate        |
| Op Erasure (`eraseOp`)                                  | Delayed           | Immediate        |
| Op Modification (`modifyOpInPlace`)                     | Immediate         | Immediate        |
| Value Replacement (`replaceAllUsesWith`)                | Delayed           | Immediate        |
| Block Insertion (`createBlock`)                         | Immediate         | Immediate        |
| Block Replacement                                       | Not supported     | Not supported    |
| Block Erasure                                           | Partly delayed    | Immediate        |
| Block Signature Conversion (`applySignatureConversion`) | Partially delayed | Immediate        |
| Region / Block Inlining  (`inlineBlockBefore`, etc.)    | Partially delayed | Immediate        |

Value replacement is delayed and has different semantics in rollback mode:
Since the actual replacement is delayed to the end of the conversion process,
additional uses of the replaced value can be created after the
`replaceAllUsesWith` call. Those uses will also be replaced at the end of the
conversion process.

Block replacement is not supported in either mode, because the rewriter
infrastructure currently has no API for replacing blocks: there is no overload
of `replaceAllUsesWith` that accepts `Block *`.

Block erasure is partly delayed in rollback mode: the block is detached from
the IR graph, but not memory for the block is not released until the end of the
conversion process. This mechanism ensures that block pointers do not change
when a block erasure is rolled back.

Block signature conversion is a combination of block insertion, op insertion,
value replacement and block erasure. In rollback mode, the first two steps are
immediate, but the last two steps are delayed.

Region / block inlining is a combination of block / op insertion and
(optionally) value replacement. In rollback mode, the insertion steps are
immediate, but the replacement step is delayed.

Note: When running in rollback mode, the conversion driver inserts fewer
transitory `builtin.unrealized_conversion_cast` ops. Such ops are needed less
frequently because certain IR modifications are delayed, making it unnecessary
to connect old (not yet rewritten) and new (already rewritten) IR in a
type-safe way. This has a negative effect on the debugging experience: when
dumping IR throughout the conversion process, users see a mixture of old and
new IR, but the way they are connected is not always visibile in the IR. Some
of that information is stored in internal C++ data structures that is not
visibile during an IR dump.

#### Type Safety

The types of the remapped operands provided to a conversion pattern (through
the adaptor or `ArrayRef` of operands) depend on type conversion rules.

If the pattern was initialized with a [type converter](#type-converter), the
conversion driver passes values whose types match the legalized types of the
operands of the matched operation as per the type converter. To that end, the
conversion driver may insert target materializations to convert the most
recently mapped values to the expected legalized types. The driver tries to
reuse existing materializations on a best-effort basis, but this is not
guaranteed by the infrastructure. If the operand types of the matched op could
not be legalized, the pattern fails to apply before the `matchAndRewrite` hook
is invoked.

Example:
```c++
// Type converter that converts all FloatTypes to IntegerTypes.
TypeConverter converter;
converter.addConversion([](FloatType t) {
    return IntegerType::get(t.getContext(), t.getWidth());
});

// Assuming that `MyConversionPattern` was initialized with `converter`.
struct MyConversionPattern : public ConversionPattern {
  virtual LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands, /* ... */) const {
//                                               ^^^^^^^^
//      If `op` has a FloatType operand, the respective value in `operands`
//      is guaranteed to have the legalized IntegerType. If another pattern
//      previously replaced the operand SSA value with an SSA value of the
//      legalized type (via "replaceOp" or "applySignatureConversion"), you
//      will get that SSA value directly (unless the replacement value was
//      also replaced). Otherwise, you will get a materialization to the
//      legalized type.
```

If the pattern was initialized without a type converter, the conversion driver
passes the most recently mapped values to the pattern, excluding any
materializations. If a value with the same type as the original operand is
desired, users can directly take the respective operand from the matched
operation.

Example: When initializing the pattern from the above example without a type
converter, `operands` contains the most recent replacement values, regardless
of their types.

Note: When running without a type converter, materializations are intentionally
excluded from the lookup process because their presence may depend on other
patterns. Passing materializations would make the conversion infrastructure
fragile and unpredictable. Moreover, there could be multiple materializations
to different types. (This can be the case when multiple patterns are running
with different type converters.) In such a case, it would be unclear which
materialization to pass.

The above rules ensure 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](#type-conversion) below.

## Type Conversion

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.

### Type Converter

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 source `Type` should be converted to N
target types. If the source type is converted to itself, we say it is a "legal"
type. Type conversions are specified via the `addConversion` method described
below.

There are two kind of conversion functions: context-aware and context-unaware
conversions. A context-unaware conversion function converts a `Type` into a
`Type`. A context-aware conversion function converts a `Value` into a type. The
latter allows users to customize type conversion rules based on the IR.

Note: context-aware type conversion functions impact the ability of the
framework to cache the conversion result. In the absence of a context-aware
conversion, all context-free type conversions can be cached. Otherwise only the
context-free conversions added after a context-aware type conversion can be
cached (conversions are applied in reverse order). 
As such it is advised to add context-aware conversions as early as possible in
the sequence of `addConversion` calls (so that they apply last).

A `materialization` describes how a list of values should be converted to a
list of values with specific types. An important distinction from 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.

*   Source Materialization

    -   A source materialization is used when a value was replaced with a value
        of a different type, but there are still users that expects the original
        ("source") type at the end of the conversion process. A source
        materialization converts the replacement value back to the source type.
    -   This materialization is used in the following situations:
        *   When a block argument has been converted to a different type, but
            the original argument still has users that will remain live after
            the conversion process has finished.
        *   When a block argument has been dropped, but the argument still has
            users that will remain live after the conversion process has
            finished.
        *   When the result type of an operation has been converted to a
            different type, but the original result still has users that will
            remain live after the conversion process is finished.

*   Target Materialization

    -   A target materialization converts a value to the type that is expected
        by a conversion pattern according to its type converter.
    -   A target materialization is used when a pattern expects the remapped
        operands to be of a certain set of types, but the original input
        operands have either not been replaced or been replaced with values of
        a different type.

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 materialization is
required, but cannot be performed, the entire conversion process fails.

Several of the available hooks are detailed below:

```c++
class TypeConverter {
 public:
  /// Register a conversion function. A conversion function must be convertible
  /// to any of the following forms (where `T` is `Value` or a class derived
  /// from `Type`, including `Type` itself):
  ///
  ///   * std::optional<Type>(T)
  ///     - This form represents a 1-1 type conversion. It should return nullptr
  ///       or `std::nullopt` to signify failure. If `std::nullopt` is returned,
  ///       the converter is allowed to try another conversion function to
  ///       perform the conversion.
  ///   * std::optional<LogicalResult>(T, SmallVectorImpl<Type> &)
  ///     - This form represents a 1-N type conversion. It should return
  ///       `failure` or `std::nullopt` 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
  ///       `std::nullopt` is returned, the converter is allowed to try another
  ///       conversion function to perform the conversion.
  ///
  /// Conversion functions that accept `Value` as the first argument are
  /// context-aware. I.e., they can take into account IR when converting the
  /// type of the given value. Context-unaware conversion functions accept
  /// `Type` or a derived class as the first argument.
  ///
  /// Note: Context-unaware conversions are cached, but context-aware
  /// conversions are not.
  ///
  /// 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)));
  }

  /// All of the following materializations require function objects that are
  /// convertible to the following form:
  ///   `std::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
  /// "casts" 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 `std::nullopt` if
  /// it failed but other materialization can be attempted, and `nullptr` on
  /// unrecoverable failure. It will only be called for (sub)types of `T`.
  /// Materialization functions must be provided when a type conversion may
  /// persist after the conversion has finished.

  /// This method registers a materialization that will be called when
  /// converting a replacement value back to its original source type.
  /// This is used when some uses of the original value 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(
        wrapSourceMaterialization<T>(std::forward<FnT>(callback)));
  }

  /// This method registers a materialization that will be called when
  /// converting a value to a target type according to a pattern's type
  /// converter.
  ///
  /// Note: Target materializations can optionally inspect the "original"
  /// type. This type may be different from the type of the input value.
  /// For example, let's assume that a conversion pattern "P1" replaced an SSA
  /// value "v1" (type "t1") with "v2" (type "t2"). Then a different conversion
  /// pattern "P2" matches an op that has "v1" as an operand. Let's furthermore
  /// assume that "P2" determines that the converted target type of "t1" is
  /// "t3", which may be different from "t2". In this example, the target
  /// materialization will be invoked with: outputType = "t3", inputs = "v2",
  /// originalType = "t1". Note that the original type "t1" cannot be recovered
  /// from just "t3" and "v2"; that's why the originalType parameter exists.
  ///
  /// Note: During a 1:N conversion, the result types can be a TypeRange. In
  /// that case the materialization produces a SmallVector<Value>.
  template <typename FnT,
            typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
  void addTargetMaterialization(FnT &&callback) {
    targetMaterializations.emplace_back(
        wrapTargetMaterialization<T>(std::forward<FnT>(callback)));
  }
};
```

Materializations through the type converter are optional. If the
`ConversionConfig::buildMaterializations` flag is set to "false", the dialect
conversion driver builds an `unrealized_conversion_cast` op instead of calling
the respective type converter callback whenever a materialization is required.

### Region Signature Conversion

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 of a
given region. 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 the operation, e.g., `func::FuncOp`, `AffineForOp`,
etc.

To convert the signature of just one given block, the
`applySignatureConversion` hook can be used.

A signature conversion, `TypeConverter::SignatureConversion`, can be built
programmatically:

```c++
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)`.

## Debugging

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 : 'func.return'(0x608000002e20) {
  "func.return"() : () -> ()

  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'func.return -> ()' {
    ** Insert  : 'spirv.Return'(0x6070000453e0)
    ** Replace : 'func.return'(0x608000002e20)

    //===-------------------------------------------===//
    Legalizing operation : 'spirv.Return'(0x6070000453e0) {
      "spirv.Return"() : () -> ()

    } -> SUCCESS : operation marked legal by the target
    //===-------------------------------------------===//
  } -> SUCCESS : pattern applied successfully
} -> SUCCESS
//===-------------------------------------------===//
```

This output is describing the legalization of an `func.return` operation. We
first try to legalize by folding the operation, but that is unsuccessful for
`func.return`. From there, a pattern is applied that replaces the `func.return`
with a `spirv.Return`. The newly generated `spirv.Return` is then processed for
legalization, but is found to already legal as per the target.