Closed stellaraccident closed 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.
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 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.
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.
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).
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.
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.
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).
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.
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
TF changes in https://github.com/iree-org/iree/pull/9804 (restriction, conceptually, from before but not in practice we believe).
Landed the final migration into iree-jax
. We are fully on the ml_program
train.
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 basedfunc
and Graph-basedsubgraph
) 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:
iree-import-public
pass to rewriteml_program.func
->func.func
(at a future point, we may just want to either useml_program.func
more deeply or have IREE have its own func).iree-import-public
conversions for globals toutil
dialect globals.TFLite:
ml_program.func
(or add a stop-gap conversion pass to make things line up for now).ml_program.global
and support externalizing the actual values.ml_program.global_*
ops.TF:
iree-jax:
ml_program.func
functions for traced wrappers (potentially producingml_program.subgraph
for HLO-only kernels).ml_program.global_*
ops instead ofiree_input.global_*
ops.torch-mlir
ml_program.func
ops instead offunc.func
.Once completed, we should no longer have:
integrations/tensorflow/iree-dialects
shadow copy.integrations/tensorflow
(we may still have a WORKSPACE that points to the TensorFlow commit we like so that we can build binaries).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.