Instead of importing a model, is it possible to give IREE as input an MLIR function?

258 views
Skip to first unread message

Do Po

unread,
Jul 8, 2021, 2:27:28 AM7/8/21
to iree-discuss

Hello,

The IREE description (https://google.github.io/iree/) says that you need to start by importing a TF model. 

But is it possible to directly provide not the model itself, but an MLIR function, like the prediction function that can be extracted from the TF model?

If yes, how would you do it? And is there a specification of what is expected of this function? I guess it must only use TF dialect operations.

Best regards,
Dumitru

Stella Laurenzo

unread,
Jul 8, 2021, 3:06:53 AM7/8/21
to Do Po, iree-discuss
The compiler accepts MHLO, Tosa or linalg containing MLIR functions. While we aim for comprehensive support of those inputs, each is presently a subset, and we are finishing them based on e2e models, primarily via the frontends for Tensorflow, TFLite, PyTorch and Jax. Aside from simple things, few are hand coding IR of any complexity to feed the compiler, and we consider this layer to be fairly mutable and are not ready to write down a firm spec for it (with the exception of tosa, which has a spec).

The core compiler is meant to be fairly hackable and may satisfy needs outside of the existing frontends but supporting them would be use case/priority driven. It isn't quite clear from your question what you are trying to do.

--
You received this message because you are subscribed to the Google Groups "iree-discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to iree-discuss...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/iree-discuss/d20f6d2f-20f3-4cc0-b2b7-ba07e82f51d1n%40googlegroups.com.

Do Po

unread,
Jul 8, 2021, 9:46:33 AM7/8/21
to iree-discuss
I'm not interested in using only the compiler (which we already tested), but also the full back-end and run-time.
Here is what I'd like to do:
- Provide an MLIR function. This function can be generated by TF or Jax, but it's important to be able to fully separate front-end and back-end at this point and know that nothing else is needed for compilation and execution.
- Be able to call the compiler and run-time on the MLIR function using a command line I control and understand. Running the resulting code should ideally not be done from Python, but from smth like C/C++, to allow embedded integration with some other code. The target architectures I'm interested in (to execute code on) in the short term are ARM and x86_64.

Dumitru

Stella Laurenzo

unread,
Jul 8, 2021, 11:09:30 AM7/8/21
to Do Po, iree-discuss
Yes, those deployment style workflows is really what IREE is aiming at itself (as opposed to online/integrated/jit, which could be built with IREE but is not a focus of ours directly). The python compiler API is just a thin wrapper around native command line tools. In fact, if you pip install it, the tools will be added to your path (ie. iree-import-tf and iree-translate). The runtime API is a) a pybind wrapper around the c API and b) some higher level python glue for interop (ie. Exposing things as python types, etc).

The python API is convenient for testing and interfacing to ml systems but deployment is intended via the underlying c API (or it's wrappers like the Android API). If you enable tracing of the python API as per here (https://github.com/google/iree/blob/main/bindings/python/iree/runtime/README.md), it will log out the sequence of c api calls needed to do whatever was needed in python (and these can be replayed by the tools iree-run-trace and iree-benchmark-trace). We use this flow in order to quickly test workloads on devices without any code being written. (Caveat: Currently these tools are not directly portable to bare metal, which will require writing the test harness code manually). Again, the python API is for convenience because so much of ml is done in python -- but it only uses public command line binaries or c APIs -- nothing special except glue).

The sample that is closest to what you are asking for is here: https://github.com/google/iree/tree/main/iree/samples/dynamic_shapes

You'll notice that we often do model extraction from python for convenience, but it is just command line tools under the covers, and when developing, we use those directly.



Reply all
Reply to author
Forward
0 new messages