Extending type conversion infrastructure

MLIR
1

@River707 and I intersected on a code review for some type conversion utilities I had written for IREE (commit) that might be good additions to the general type conversion infra in MLIR. I’m restarting that discussion here so we don’t lose it and can continue it in the right forum.

— Comment from @River707

There are quite a few commonalities between this and the TypeConverter upstream. Are there ways that we can improve that and reduce the amount of code duplication? Type conversion is something that hasn’t been heavily invested in, so it’d be nice to improve things as much as possible upstream instead of doing it out-of-tree. (Even if it isn’t directly applicable to applyConversionPatterns)

— Response from @stellaraccident

Agreed and I would love to investigate that. Indeed, before me accidentally including a debug change in a core file caused you to be CC’d, I had intended to land this, make sure it is what I wanted with a bit more incremental work and then talk to you about it. This came after quite a bit of time spent trying to get the upstream TypeConverter bits to work, hitting the limits of the abstraction and hand-coding it. After I hand-coded three variants of similar things, I boiled out this facility.

The two things I found missing from the upstream TypeConverter were:

  • Casts both ways. Upstream only has the equivalent of castToTarget and I was finding myself creating a fair bit of grungy code to bridge that.
  • Ability to use TypeConverters to make local changes to the IR outside of the conversion framework. (might be my limitations in grokking it but I found sharp edges that I couldn’t overcome).

Ergonomically, I found a few other things:

  • The upstream conversion-framework based TypeConverter does a reasonable job for simple 1:1 conversions but has missing features for more advanced cases (or they were beyond my ability to figure out). Compare for example, the code deleted in ExpandFunctionDynamicDims. There was an equal bit of complex code in OutlineDispatchRegions before I rewrote it to this form.
  • After having written the signature rewriting code 2-3 times with TypeConverter, I was still running into a lot of trial and error “holding it wrong” cases each time. I think a lot of these come from the way that the conversion framework needs to process things, and more general passes don’t need the complexity (and ime, tend to be doing weirder things anyway and need the flexibility to make consistent local modifications).

I’d like to figure out some specific ways to contribute some of this functionality upstream. I might be getting hung up on trivialities: the current TypeConverter is embedded/entwined with the DialectConversion framework in a way that I’m not sure layers how I would want. In my mind, there might be a lower level facility which:

  • allows me to define a type conversion in terms of the ability to convert types from :
    • source->target*
    • materialize a cast from source->target* values
    • materialize a cast from target*->source values
  • this would be sufficient for the generic construction of local conversion utilities for rewriting:
    • function signatures
    • function calls
    • block signatures
    • terminators
    • generic value lists

Currently the upstream TypeConverter has just enough defined to convert the entry block signature in a function signature conversion. Also, having used it a few times, I still find the actual process of using it to transform function signatures to be somewhat verbose and error prone (and tends to require special case fixups to terminators/returns/etc). It seems like if, instead of just taking a lambda Type->SmallVector<Type>, if it instead took a full converter like my TypeExpander it could be made a lot more general. Also, with that more general facility, I’ve found it easier to just apply local type conversions to things like function signatures, block signatures and calls independently versus doing the accounting that SignatureConversion does.

I wouldn’t call what I’m doing a fully generic type converter (there are plenty of other ways that type conversion needs to happen in practice), but I have found that this pattern of type expansion happens semi-frequently as you lower from higher level dialects, and having it wrapped up as I have simplified a lot of code.

Advice?

2

My advice would be to do just that. We could change the internal storage to be a set of TypeExpander objects instead of callbacks. The existing lambda methods can just construct a simple TypeExpander internally. That way users that want the more general functionality can get it, and no existing users have to be change. This is the direction that I intended to move towards when I made the lambda callbacks in the first place. The only thing I would be against is using inheritance for extension, ideally all of the virtual methods on TypeConverter go away and users just provide these “extension”(TypeExpander) objects to control behavior.

An easy first step to break the direct connection with dialect conversion would be to move TypeConverter to a new header file TypeConversion. At that point we can establish that dialect conversion is one (well main) user, but not the only one.

3

Thanks - this sounds tractable and I can take a stab at it this weekend. Could you elaborate a bit on the above statement, though? Would there be one “TypeExpander” that takes lambdas or still a base class that only had virtual methods on it (or some other template based specialization pattern). I think I agree with the sentiment but just don’t have the experience you do wrt polymorphism without virtual methods and could use some more specific tips.

4

Basically what I want to avoid is having the main “TypeConverter” class have virtual methods. At least for dialect conversion, the conversion related objects(conversion target, type converter, etc.) end up being passed around and initialized by various different sources(depending on which conversion you add). It effectively breaks the composability of conversions. What is fine though, is letting the extension class use virtual methods; e.g. something like:

struct TypeConverter {
  /// The class users provide for extending the converter.
  struct Extension {
    virtual LogicalResult convertType(...);
    ...
  };
  void registerExtension(std::unique_ptr<Extension>);

  /// Mapping the existing API.
  void addConversion(ConversionCallbackFn callback) {
    struct SimpleExtension : public Extension {
        SimpleExtension(ConversionCallbackFn callback) : callback(callback) {}
        LogicalResult convertType(...) override { return callback(...); }

       ConversionCallbackFn callback;
    };
   registerExtension(std::make_unique<SimpleExtension>(callback));
  }
};

Does that make sense?

1 Like
5

Perfect - that clears it up for me! Thanks!

6

Hey @River707, I took a stab at this but got kind of stuck trying to figure out what level of surgery to do on the dialect-conversion side – and decided to ask for more detail on what your desired end state is before trying to force it.

Here’s what I’m working on: Add TypeConversionPattern class · llvm/llvm-project@10108af · GitHub

For reference, here’s the facility in IREE that I’m basing this on (I’ve used it pretty extensively now in favor over the existing type conversion infra and have experienced fewer sharp edges): https://github.com/google/iree/blob/master/iree/compiler/Dialect/Shape/Utils/TypeConversion.h

I suspect that if I mentally factor the existing SignatureConversion stuff out of the TypeConverter and look at ways to simplify that separately, I’m going to have a better time. Since I don’t know why the current signature conversion needs to do all of the accounting it does (and have a correspondingly large public surface area), I’m having trouble knowing what it is ok to just strip down. I’ve basically implemented the same facility in a more general way (I think) with just a couple of public functions, but I may be missing a subtlety.

There aren’t that many users of the existing type conversion infra yet. I guess my question is how much would it be better to define a lower-level facility outside of dialect-conversion, then implement some simpler ways to do dialect-conversion in terms of that, and then migrate call sites to the newer form? Since I’m having trouble with the complexity of the existing type conversion paths, I’m having trouble figuring out how to factor a layer out of them in-situ.

7

SignatureConversion is essentially a mechanism in which to describe the result of converting the signature of a block. It includes information about how types get remapped as DialectConversion needs to be able to remap the existing arguments to the new ones. Remember that DialectConversion doesn’t perform replace values directly, and needs to maintain a mapping between old and new values. I logically separate things into 3 distinct layers:

DialectConversion:
The executor of the type conversion. This is where casts get materialized, blocks get their arguments changed, etc.

SignatureConversion:
This is the logical plan of the conversion that is provided to DialectConversion.

TypeConversion:
This is the policy object, which defines how a specific type gets converted. It provides ways to materialize casts back-and-forth if necessary, but generally only focuses on converting existing types.

I separate SignatureConversion from the TypeConversion for a couple of reasons:

  • Types may be added out of nowhere

    • Depending on the conversion, a type can seemingly be materialized out of nowhere. This isn’t something that a policy object like TypeConversion can, or should handle.

  • For conversions specifically, we may want to replace an BlockArgument with a different Value

    • This is also something that a policy object shouldn’t really handle. This type of conversion is context specific depending on the specific Pattern being applied, and isn’t a general rule that can be applied. It is a special case of a 1->0 mapping.

  • The mapping of some types is context dependent

    • Some signatures may be converted different contextually depending on some properties of the value. For example, the argument of a function may be converted differently depending on the presence of an attribute of its uses within the function.

To summarize, SignatureConversion was intended as a mechanism to allow building a plan for a conversion that takes into account the special contextual circumstances that popup during conversion. Not sure if that makes sense, but that was the original rationale.

8

This is a roundabout way of saying that I’m fine with inverting the dependency chain such that SignatureConversion depends on TypeConverter instead of the other way around. They were both only ever used by DialectConversion so they grew together.

9

Ok, thank you – that helps me understand the layering better. I was confusing myself due to the nesting of SignatureConversion (and related things) in TypeConverter. The mechanics of how dialect conversion actually works are still fuzzy to me, but I don’t need to understand all of the details.

So, in terms of end state, I imagine that (as discussed previously):

  • SignatureConversion: Gets promoted to a top-level type
  • TypeConverter becomes just the policy object with some methods moving to SignatureConversion (and moving to a new top-level header)
    • convertSignatureArg
    • convertBlockSignature
    • (possibly) isSignatureLegal
  • TypeConverter::materializeConversion becomes something like TypeConverter::materializeTargetConversion (or different spelling) and returns a Value, allowing it to handle identity and 1:1 cases better
  • New TypeConverter::materializeSourceConversion (or different spelling)
  • May also add some non-conversion based in-place type update helpers to TypeConversion.h (similar to what I have in IREE’s TypeExpander). These have been useful when operating outside of the conversion framework, and some transformations are really tricky (esp. arg/result attributes and such), and would be good to centralize.

In terms of getting there, I’ll probably extract SignatureConversion and its new methods first.

Make sense?

10

Sounds good to me.

[Disregard, for some reason I need 20 characters to post]