Executing MLIR pipelines via pip installed TensorFlow

159 views
Skip to first unread message

Brehler, Marius

unread,
Oct 5, 2021, 4:55:06 PM10/5/21
to ml...@tensorflow.org
Hi all,

To end2end test EmitC, I had set up a GitHub Action based pipeline that converts a model all the way down to C++. This pipeline used a pip installed TensorFlow version to avoid building tf-opt in the GiHub Action.
The input is a SavedModel that is translated into TF dialect in the first step. Next, TF is converted to MHLO. Finally, this is converted to EmitC and translated to C++.
While using TensorFlow 2.5, we could do this by utilizing `pywrap_mlir.experimental_run_pass_pipeline()`. This was possible as long as passes like `xla-legalize-tf` were registered statically (https://github.com/tensorflow/tensorflow/blob/8222c1cfc866126111f23bd9872998480cebf2c1/tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc#L6290-L6291). We could define a pipeline like

pass_pipeline = ",".join([
    "func(", "canonicalize",
    "tf-saved-model-optimize-global-tensors"
])

and pass this to `pywrap_mlir.experimental_run_pass_pipeline()`. For some passes the static registration was replaced by a declarative one (e.g. for `xla-legalize-tf` with https://github.com/tensorflow/tensorflow/commit/bdafd9a4e4c26a852be79ce0153d84f513f3525d). With these changes, it seems that our scripts cannot be used that way any longer:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Invalid pass_pipeline: MLIR Textual PassPipeline Parser:1:6: error: 'xla-legalize-tf' does not refer to a registered pass or pass pipeline

I am aware that we're using an unstable API, but maybe one of you has an idea and can tell if there is still a way to run a pipeline via an pip installed python (by using `pywrap_mlir.experimental_run_pass_pipeline()` or any other function). I would really like to avoid `tf-opt` in a GH action.
Thanks in advance.

Best, Marius
-- 
Dr.-Ing. Marius Brehler
Senior Scientist
AI and Autonomous Systems

Fraunhofer Institute for Material Flow and Logistics IML
Joseph-von-Fraunhofer-Str. 2-4
44227 Dortmund, Germany

Jacques Pienaar

unread,
Oct 5, 2021, 5:04:48 PM10/5/21
to Brehler, Marius, Mehdi Amini, ml...@tensorflow.org
Hey,

@Mehdi Amini would know these best. But I think we simply haven't included the static registry for these in the existing build (basically where previously there was

static PassRegistration<LegalizeTF> pass(
    "xla-legalize-tf", "Legalize from TensorFlow to the XLA dialect");

we could know use the registration generated by ODS but haven't). I don't know if we expose that in a different way ... And this seems like it was an oversight when I made the change as we explicitly call the register calls in tf-opt. I can look at that.

-- Jacques

--
You received this message because you are subscribed to the Google Groups "MLIR" group.
To unsubscribe from this group and stop receiving emails from it, send an email to mlir+uns...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/ba14df94eae92d564b50843e74fbfb37c354d5d8.camel%40iml.fraunhofer.de.

Brehler, Marius

unread,
Oct 5, 2021, 5:22:36 PM10/5/21
to ami...@google.com, jpie...@google.com, ml...@tensorflow.org
Hey Jacques,

Thanks for the quick reply. It is not only an issue with `xla-legalize`. It is the same with `tf-device-decompose-resource-ops` (changed with https://github.com/tensorflow/tensorflow/commit/43a65c6275cc9076ef9ba224e7111c526f09ba43#diff-80d3447e0b9a264568367091e21f77e1f202c5154c489271c7185138a4c1bfa9). Therefore, it seems to be a more general problem, at least when assuming that `pywrap_mlir.experimental_run_pass_pipeline()` should be able to run these passes. I really appreciate that you're offering to take a look and I would be happy to assist, but wanted to make sure this isn't an intended change before starting to take a deeper look.

Best, Marius

Retford, Brian

unread,
Oct 5, 2021, 6:13:28 PM10/5/21
to Brehler, Marius, ami...@google.com, jpie...@google.com, ml...@tensorflow.org

We have a tiny fork of tensorflow whose sole purpose is to include all of the passes in the registry – if they found their way into the mainline pip that would be a small dream come true

Brehler, Marius

unread,
Oct 11, 2021, 9:27:54 AM10/11/21
to brian....@intel.com, ami...@google.com, jpie...@google.com, ml...@tensorflow.org
Hey Brain, hey Jacques,

@Brain: Does the fork enable to use the passes via `pywrap_mlir.experimental_run_pass_pipeline()`? Further, is it public?

@Jacques: At the moment this blocks upgrading the mlir-mhlo submodule to the current HEAD in the mlir-emitc repo. With the attributes of mhlo. convolution updated with https://github.com/tensorflow/mlir-hlo/commit/ac4d5e4, the EmitC conversions need an update. Also the mlir-emtic repo no longer includes the EmitC dialect, it still demonstrates how to use EmitC on the basis of an end2end test. As described before this is now broken. Since you mentioned that you will take a look: Is there n agreement of the TF side that the passes should be available `pywrap_mlir.experimental_run_pass_pipeline()`? I am asking to get an idea if we can expect that those will be available via the pip installable tf-nightly in near future.
As background information: I would like to propose migrating the MHLO -> EmitC -> C++ from the mlir-emitc repo to the tensorflow repo (and I also plan to upstream the TOSA -> EmitC -> C++ pass if there is interest). However, I would first like to fix our pipeline.

Best, Marius

Jacques Pienaar

unread,
Oct 11, 2021, 9:51:31 AM10/11/21
to Brehler, Marius, Retford, Brian, Mehdi Amini, MLIR
Hey Marius,

Yes, Mehdi has a WIP change to enable that. And you know I'm pro the move proposal.

Best,

Jacques

Brehler, Marius

unread,
Oct 11, 2021, 9:57:52 AM10/11/21
to jpie...@google.com, ami...@google.com, ml...@tensorflow.org
Hey Jacques,

Awesome, thanks for the update! This information helps me a lot as I don't need to explore other potential solutions. And yes, I am aware that you're pro the move proposal :)

Best, Marius

Mehdi Amini

unread,
Oct 11, 2021, 12:14:46 PM10/11/21
to Jacques Pienaar, Brehler, Marius, Retford, Brian, MLIR
Yeah I just stuck on some issue with the annoying OSS shared library setup in TensorFlow for now, so no ETA on the fix right now.

Mehdi Amini

unread,
Oct 14, 2021, 10:59:27 PM10/14/21
to Jacques Pienaar, Brehler, Marius, Retford, Brian, MLIR
I landed a few things, can you try again at head and let me know?

Brehler, Marius

unread,
Oct 18, 2021, 6:28:35 AM10/18/21
to ami...@google.com, jpie...@google.com, ml...@tensorflow.org
Hey Mehdi,
Thanks for your efforts!

I was able to run the passes `xla-legalize` and `tf-device-decompose-resource-ops` via `pywrap_mlir.experimental_run_pass_pipeline()` with `tf-nightly==2.8.0.dev20211017`.
Unfortunately, `tf-saved-model-optimize-global-tensors` and `tf-saved-model-freeze-global-tensors`, both defined in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/transforms/tf_savedmodel_passes.td, seem to be unregistered as they fail with "tensorflow.python.framework.errors_impl.InvalidArgumentError: Invalid pass_pipeline: MLIR Textual PassPipeline Parser:1:218: error: 'tf-saved-model-{optimize-global-tensor,freeze-global-tensors}' does not refer to a registered pass or pass pipeline".

Best, Marius

Brehler, Marius

unread,
Oct 20, 2021, 10:00:22 AM10/20/21
to ami...@google.com, jpie...@google.com, ml...@tensorflow.org
I recently noticed that also `tf-to-tosa-pipeline` isn't a registered pipeline when using `pywrap_mlir.experimental_run_pass_pipeline()`.

Mehdi AMINI

unread,
Nov 2, 2021, 1:31:36 AM11/2/21
to Brehler, Marius, ami...@google.com, jpie...@google.com, ml...@tensorflow.org
Hey,

I had to land a bunch of refactoring (14 commits), but it should all be there now. Let me know if anything else is missing!

Cheers,

-- 
Mehdi


Brehler, Marius

unread,
Nov 2, 2021, 12:25:29 PM11/2/21
to joke...@gmail.com, ami...@google.com, ml...@tensorflow.org, jpie...@google.com
Hey Mehdi,

I tested with tf-nightly==2.8.0.dev20211102 a minute ago and all passes we're using seem to work as expected. Thanks a lot!

Best, Marius
Reply all
Reply to author
Forward
0 new messages