Understanding type conversions during dialect conversions

MLIR
1

Hi!
I’m trying to get a detailed understanding of dialect conversions, in particular of how to handle type conversions and conversions of operations from unrelated dialects. What is unclear to me is:

  • The relationship of source, target and argument materializations to legalization (e.g., via rewrite patterns)
  • The conversion of polymorphic operations neither belonging to the source nor the target dialect

As an example for a conversion scenario, consider the following function using FooDialect that should be converted to operations and types from TargetDialect. The function below:

func.func @foo(%arg0: !FooDialect.foo_type, %arg1: !FooDialect.foo_type) -> !FooDialect.foo_type
{
   %0 = FooDialect.some_op %arg0, %arg1 : (!FooDialect.foo_type, !FooDialect.foo_type) -> !FooDialect.foo_type
   return %0 : !FooDialect.foo_type
}

should become:

func.func @foo(%arg0: !TargetDialect.target_type, %arg1: !TargetDialect.target_type) -> !TargetDialect.target_type
{
   %0 = TargetDialect.another_op %arg0, %arg1 : (!TargetDialect.target_type, !TargetDialect.target_type) -> !TargetDialect.target_type
   return %0 : !TargetDialect.target_type
}

As far as I understand, the following is needed to implement the conversion:

  • A conversion target specification marking:
    • the FooDialect as illegal
    • func::ReturnOp dynamically as illegal if the operand and result types belong to FooDialect
  • A type converter with:
    • an argument materialization creating builtin.unrealized_conversion_casts
    • a source materialization creating builtin.unrealized_conversion_casts
    • a target materialization creating builtin.unrealized_conversion_casts
  • A ConversionPatternSet with
    • A pattern for each operation from FooDialect, replacing the op with an op from TargetDialect
    • An instance of the FunctionOpInterfaceSignatureConversion pattern that converts the function type and updates the function in place
    • A pattern for func::ReturnOp

Furthermore, I understand that the unrealized_conversion_casts are only there to temporarily reconcile the IR until all operations have been converted.

I assume that source, target and argument conversions are triggered whenever an illegal op is processed, whose operands or results have a type for which a respective materialization exists. I further assume that this is completely transparent to the legalization pattern itself (e.g., a rewrite pattern), as operands are rewritten during materialization and before the illegal op is finally processed. Is this correct?

Let’s go through the conversion step by step. The first thing that should happen is the signature conversion. Since there are still consumers of the arguments with the old type, signature conversion should kick in and insert builtin.unrealized_conversion_casts for each argument, leaving the IR as:

func.func @foo(%arg0: !TargetDialect.target_type, %arg1: !TargetDialect.target_type) -> !TargetDialect.target_type
{
  %tmpArg0 = builtin.unrealized_conversion_cast %arg0 : !TargetDialect.target_type to !FooDialect.foo_type
  %tmpArg1 = builtin.unrealized_conversion_cast %arg1 : !TargetDialect.target_type to !FooDialect.foo_type
   %0 = FooDialect.some_op %tmpArg0, %tmpArg1 : (!FooDialect.foo_type, !FooDialect.foo_type) -> !FooDialect.foo_type
   return %0 : !FooDialect.foo_type
}

I’m not sure about how the return operation is handled. Is the above state, the IR is invalid, since the operand type of the return operation does not match the return type of the function operation. Supposing that the IR remains just invalid and the conversion attempts to reconcile it, the pattern for the return op should be executed next, triggering source materialization and resulting in:

func.func @foo(%arg0: !TargetDialect.target_type, %arg1: !TargetDialect.target_type) -> !TargetDialect.target_type
{
   %tmpArg0 = builtin.unrealized_conversion_cast %arg0 : !TargetDialect.target_type to !FooDialect.foo_type
   %tmpArg1 = builtin.unrealized_conversion_cast %arg1 : !TargetDialect.target_type to !FooDialect.foo_type
   %0 = FooDialect.some_op %tmpArg0, %tmpArg1 : (!FooDialect.foo_type, !FooDialect.foo_type) -> !FooDialect.foo_type
   %0converted = builtin.unrealized_conversion_cast %0 : !FooDialect.foo_type to !TargetDialect.target_type
   return %0converted : !TargetDialect.target_type
}

Finally, the pattern for FooDialect.some_op would be executed, triggering source and target materialization:

func.func @foo(%arg0: !TargetDialect.target_type, %arg1: !TargetDialect.target_type) -> !TargetDialect.target_type
{
   %tmpArg0 = builtin.unrealized_conversion_cast %arg0 : !TargetDialect.target_type to !FooDialect.foo_type
   %tmpArg1 = builtin.unrealized_conversion_cast %arg1 : !TargetDialect.target_type to !FooDialect.foo_type
   %tmpArgReconverted0 = builtin.unrealized_conversion_cast %tmpArg0 : !FooDialect.foo_type to  !TargetDialect.target_type
   %tmpArgReconverted1 = builtin.unrealized_conversion_cast %tmpArg1 : !FooDialect.foo_type to  !TargetDialect.target_type
   %0 = TargetDialect.target_op %tmpArgReconverted0, %tmpArgReconverted1 : (!TargetDialect.target_type, !TargetDialect.target_type) -> !TargetDialect.target_type
   %tmp0 = builtin.unrealized_conversion_cast %arg0 : !TargetDialect.target_type to !FooDialect.foo_type
   %0converted = builtin.unrealized_conversion_cast %tmp0 : !FooDialect.foo_type to !TargetDialect.target_type
   return %0converted : !TargetDialect.target_type
}

And then finally, the unrealized_conversion_casts are folded away:

func.func @foo(%arg0: !TargetDialect.target_type, %arg1: !TargetDialect.target_type) -> !TargetDialect.target_type
{
   %0 = TargetDialect.target_op %arg0, %arg1 : (!TargetDialect.target_type, !TargetDialect.target_type) -> !TargetDialect.target_type
   return %0 : !TargetDialect.target_type
}

Besides the handling of the return op after the signature conversion, I also wonder if there is a generic pattern that can be applied to handle polymorphic operations from unrelated dialects in general. At least, it seems odd to me to have a pattern for func::ReturnOp in a conversion of FooDialect to TargetDialect. Following the same logic, there would have to be a pattern for linalg.generic, too if for some reason the IR contains linalg.generic operations working on FooDialect.foo_types.

I’d be grateful if someone could shed some light on this.

Thanks,
Andi

2

From what I’ve experienced till now, for the func.return op you can use the following pattern provided in mlir/Dialect/Func/Transforms/FuncConversions.h.

populateReturnOpTypeConversionPattern(patterns, converter);
    target.addDynamicallyLegalOp<func::ReturnOp>(
        [&](func::ReturnOp op) { return converter.isLegal(op); });

As for the linalg.generic I’m also really confused how to deal with the type conversion within the op.
BTW to see the actuall stpes of lowering you mentioned above, I normally use -debug-only=dialect-conversion to check how it works in the command line.