Executing MLIR pipelines via pip installed TensorFlow

[Brehler, Marius]

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

Mail: marius....@iml.fraunhofer.de

Web: http://www.iml.fraunhofer.de

[Jacques Pienaar]

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]

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]

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

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/110adc2efa835fadabbee63fc4b7bfce48364c03.camel%40iml.fraunhofer.de.

[Brehler, Marius]

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]

Hey Marius,

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

Best,

Jacques

[Brehler, Marius]

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]

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]

I landed a few things, can you try again at head and let me know?

[Brehler, Marius]

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]

I recently noticed that also `tf-to-tosa-pipeline` isn't a registered pipeline when using `pywrap_mlir.experimental_run_pass_pipeline()`.

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/16716e52234966730e5f8b1d0bcdf725ccd4c0f3.camel%40iml.fraunhofer.de.

[Mehdi AMINI]

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

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/8fc7de414ec05cba11e9b8e45159a20b2f0d8aef.camel%40iml.fraunhofer.de.

[Brehler, Marius]

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

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/16716e52234966730e5f8b1d0bcdf725ccd4c0f3.camel%40iml.fraunhofer.de.

[Abhishek]

Hey Marius

I try the following code to transform from tf_exec to hlo but I see an assertion failure. 

pass = 'xla-legalize-tf'

mlir_mlho_output = experimental_run_pass_pipeline(mlir, pass, true)

Error :

llvm::cast(From*) [with To = mlir::func::FuncOp; From = mlir::Operation]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.

Aborted (core dumped).

Am I missing something?

thanks

Abhishek

[Christian Sigg]

I think this should be fixed with https://github.com/tensorflow/tensorflow/commit/9a493a78daeccb9e24af2774c0898f6ef0554f9b.