iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.83k stars 611 forks source link

Untangle type conversion from mhlo->linalg pass. #10897

Open benvanik opened 2 years ago

benvanik commented 2 years ago

The MHLOToLinalgOnTensors pass is also converting type signedness inside of it and shouldn't be - we don't just have mhlo ops turning into linalg ops, but any number of custom ops or other dialect ops. The first issue we hit is #10834 and I was able to work around it in #10896 but looking at the IR I suspect there will be others that get much trickier - especially along the ABI. Converting types in generic IR is a big transformation and something that should live in its own pass so that its possible to debug and support more robust conversion (module scope, custom ops, etc).

I think there's likely some larger reworking to be done around the whole input pipeline. For example today we run the mhlo/tosa/etc pipelines before ABI generation which means we can't capture the original program types as by the time we can they've already been changed. If the type is actually changing this is ok (like our f64->f32 demotion pass), but in these signedness cases the user is still going to expect ui32 if that's what they passed in and not have them reflected as just i32. We'll get issues from users talking about ndarray dtypes not matching and such (I think we already have but punted on them).

For type conversion we should use the existing ConvertPrimitiveType helper (compiler/src/iree/compiler/Dialect/Util/Transforms/ConvertPrimitiveType.cpp) - that handles all ops and would be something we could run in the MHLO pipeline where appropriate. The conversion is tricky to get right and it'd be nice to have one fewer slightly busted implementation.

MaheshRavishankar commented 2 years ago

Just assigning to Geoffrey and Rob since they have looked at it, but probably a P2.

GMNGeoffrey commented 2 years ago

The issue is that MHLO only operates on signfull types and linalg only operates on signless types. That conversion has to happen somewhere. I believe that at the time I argued for doing this in a more staged fashion, but it would require some ops that agree to actually work with both signfull and signless types and upstream has historically refused to allow signfullness to creep into their perfect signless world. One thing that would at least help make things a bit more understandable and less error prone would be explicit sign-casting op(s) that go from signless <-> signfull. I think my preference would be separate ops for the 4 different transformations, but 2 ops or a single op would probably also work. Right now we're just using the unrealized conversion cast, I think, which means introspecting on what it's converting between any time you try to mess with it. Even with those ops though, with the current (as of like a year ago when I last looked) mechanisms for type conversion, I think we need to do this as part of converting to ops that operate on signless types. Linalg is currently the place that happens. We could imagine an opset at the same level as MHLO with the only difference being that the sign information is carried on the op instead of on the type. That actually seems the nicest to me, except that it would be annoying to maintain. Some kind of generation mechanism might help with that. I don't know how this factors into these other input dialects though.

The phase-ordering issue seems like a separate bad thing.

The MHLOToLinalgOnTensors pass is also converting type signedness inside of it and shouldn't be - we don't just have mhlo ops turning into linalg ops, but any number of custom ops or other dialect ops.

I don't quite understand. Do you mean that MHLOToLinalgOnTensors is converting things that aren't MHLO, that it's converting to things other than linalg, or just that its input could contain non-MHLO ops that it shouldn't be touching?

GMNGeoffrey commented 1 year ago

Unassigning myself from issues that I'm not actively working on