Start upgrading to `ml_program`

stellaraccident commented 2 years ago

Background: We are lobbying to introduce a new ml_program dialect upstream (initial RFC discussion). This effort aims to get sufficient operations upstream such that we don't need iree_input for most "normal" ML programs from TFLite, JAX, PyTorch. The initial scope includes container ops (CFG based func and Graph-based subgraph) and global/extern-constant infrastructure.

I would like IREE to be the first adopter of what we produce here because it will allow us to sever the dependency of the frontends on IREE dialects, at least for a wide range of cases. This will let us propose that our importers be directly included in the various frontend projects without taking a dependency on IREE (unless if something advanced is warranted). Given that we will also be building out facilities upstream for externalized constants (in sidecar files outside of the IR), once even the initial scope is done, it should also yield a pretty nice usability improvement in addition to the layering improvements.

Rough work to be done (dependent on upstream consensus/landing):

IREE:

[ ] Extend the iree-import-public pass to rewrite ml_program.func -> func.func (at a future point, we may just want to either use ml_program.func more deeply or have IREE have its own func).
[ ] Teach IREE how to handle externalized constant files if provided (this should extend the APIs to accept constant files and implement core support for using them in IREE's globals).
- [ ] Thread constant files through more pervasively, supporting generating them in global simplification/numeric optimization and potentially preserving them through to vmfb generation (there is a lot that can be done here, and fully working out is likely a project in itself -- as part of this effort, we will take some naive steps to make it work at all).
[ ] Add iree-import-public conversions for globals to util dialect globals.

TFLite:

[ ] Rework the upstream TF code to produce ml_program.func (or add a stop-gap conversion pass to make things line up for now).
[ ] Hoist large inline constants and globals to ml_program.global and support externalizing the actual values.
[ ] Update stateful op lowerings to use ml_program.global_* ops.
[ ] Contribute the TFLite importer to TensorFlow so that anyone can build the binary and get ml_program + TOSA IR.
[ ] Delete the TFLite importer from IREE.

TF:

[x] #8832

iree-jax:

[ ] Update the iree-jax Python code to produce ml_program.func functions for traced wrappers (potentially producing ml_program.subgraph for HLO-only kernels).
[ ] Emit ml_program.global_* ops instead of iree_input.global_* ops.
[ ] Emit externalized values for globals and hoist large inline constants to externs.
[ ] Raise an RFC to the JAX team to either incorporate iree-jax or rewrite it to their wishes.
[ ] Delete iree-jax.

torch-mlir

[ ] Teach torch-mlir to produce ml_program.func ops instead of func.func.
[ ] Rework value importing in torch-mlir to emit globals and allow generation of externalized data for large tensor values.
[ ] Map mutable module state to globals.

Once completed, we should no longer have:

integrations/tensorflow/iree-dialects shadow copy.
Any code in integrations/tensorflow (we may still have a WORKSPACE that points to the TensorFlow commit we like so that we can build binaries).
A standalone iree-jax project.

We will likely continue to build the TFLite importer (hopefully from pristine upstream sources) and distribute it in our wheels for usability. However, depending on how ambitious we are, it may be possible to simply expose a TF Python API to do such conversions.

silvasean commented 2 years ago

@sjarus FYI - Torch-MLIR will probably start producing ml_program.func

Extend the iree-import-public pass to rewrite ml_program.func -> func.func (at a future point, we may just want to either use ml_program.func more deeply or have IREE have its own func).

Can this just be an upstream pass? Seems like "everybody" is going to need this, at least at the beginning.

jpienaar commented 2 years ago

Btw working with Jeff and Mehdi on externalized constants. It's just an attribute so any op that uses ElementsAttr (e.g., constant) would be able to reference extern constant in a self-contained/repro producing manner. Adding additional metadata for origin files etc can be added

sjarus commented 2 years ago

@sjarus FYI - Torch-MLIR will probably start producing ml_program.func

Extend the iree-import-public pass to rewrite ml_program.func -> func.func (at a future point, we may just want to either use ml_program.func more deeply or have IREE have its own func).

Can this just be an upstream pass? Seems like "everybody" is going to need this, at least at the beginning.

I'd just sent @stellaraccident some notes yesterday on TOSA and its intersection with this construct. Sounds like the timing was pretty good. Happy to keep working in this direction.

silvasean commented 2 years ago

Btw working with Jeff and Mehdi on externalized constants. It's just an attribute so any op that uses ElementsAttr (e.g., constant) would be able to reference extern constant in a self-contained/repro producing manner. Adding additional metadata for origin files etc can be added

What happened to having it just be linked to the sym_name of the global? I thought that's what we were going for.

jpienaar commented 2 years ago

Btw working with Jeff and Mehdi on externalized constants. It's just an attribute so any op that uses ElementsAttr (e.g., constant) would be able to reference extern constant in a self-contained/repro producing manner. Adding additional metadata for origin files etc can be added

What happened to having it just be linked to the sym_name of the global? I thought that's what we were going for.

Sym_name doesn't avoid uniqueing in context unless you make the IR non-self contained (which works for cases where you just want constant opaque, or where you aren't running passes as part of another system but as standalone tool).

silvasean commented 2 years ago

Btw working with Jeff and Mehdi on externalized constants. It's just an attribute so any op that uses ElementsAttr (e.g., constant) would be able to reference extern constant in a self-contained/repro producing manner. Adding additional metadata for origin files etc can be added

What happened to having it just be linked to the sym_name of the global? I thought that's what we were going for.

Sym_name doesn't avoid uniqueing in context unless you make the IR non-self contained (which works for cases where you just want constant opaque, or where you aren't running passes as part of another system but as standalone tool).

Let's sync up offline.

stellaraccident commented 2 years ago

Sym_name doesn't avoid uniqueing in context unless you make the IR non-self contained (which works for cases where you just want constant opaque, or where you aren't running passes as part of another system but as standalone tool).

Let's sync up offline.

Let's chat more about this. The approach described in the RFC was actually (imo) a really nice simplification from Chris that will play really well with the tools. That doesn't need to be the only way to do it, but extern binding by name matches most other programming models, has strong priors, etc.

jpienaar commented 2 years ago

Sym_name doesn't avoid uniqueing in context unless you make the IR non-self contained (which works for cases where you just want constant opaque, or where you aren't running passes as part of another system but as standalone tool).

Let's sync up offline.

Let's chat more about this. The approach described in the RFC was actually (imo) a really nice simplification from Chris that will play really well with the tools. That doesn't need to be the only way to do it, but extern binding by name matches most other programming models, has strong priors, etc.

Sounds good, note that it is orthogonal to binding by name or using global ops.

I'll start roping some folks in for TFlite and JAX work. For Torch I'm assuming Sean is on it (but correct me if wrong :-)).

Hoist large inline constants and globals to ml_program.global and support externalizing the actual values.

Tflite export does externalization at the moment when going to flatbuffer, what externalization options will the other backends be using? (I'm trying to think if we should use TOSA's flatbuffer or not and so have the result be a relatively pure TOSA export).

Either delete the TF importer or upgrade it in a similar fashion to the above. #8832

This may be a good question for bridge team too: seems much of the functionality here is bridge related and could build on APIs they provide/most parts moved there.

Would the minimal here be to produce ml_program.func (unless HLO only) instead of func.func? (Seems with that and if we only care about MHLO section, we could reuse parts for Jax in terms of hoisting and this be reasonably small in scope).

sjarus commented 2 years ago

Tflite export does externalization at the moment when going to flatbuffer, what externalization options will the other backends be using? (I'm trying to think if we should use TOSA's flatbuffer or not and so have the result be a relatively pure TOSA export).

The TOSA serialization format has certain design choices that have evolved over time, e.g. the choice of embedding weight content vs reference to a .npy file containing weights. If you're interested in using the format, we'd be happy to get you connected to conversations internally so you're up to date on our thinking and can offer your own feedback.

rsuderman commented 2 years ago

IREE specific changes have landed to support ml_program so we can start shifting front-ends over.

JAX changes for this are lined up in https://github.com/rsuderman/jax/commit/3c3c96cb905ab3f41106cb6dc0b148e778e7d665

Logged a separate issue for switching tflite over to ml_program https://github.com/iree-org/iree/issues/9654

jpienaar commented 2 years ago

TF changes in https://github.com/iree-org/iree/pull/9804 (restriction, conceptually, from before but not in practice we believe).

rsuderman commented 2 years ago

Landed the final migration into iree-jax. We are fully on the ml_program train.

https://github.com/iree-org/iree-jax/pull/23

iree-org / iree