TensorFlow dialect availability and compatibility


Is there a TensorFlow dialect available for MLIR?

A quick search gives this: MLIR dialects  |  TensorFlow MLIR
But the page seems outdated.
Furthermore, if I look on GitHub - tensorflow/mlir: "Multi-Level Intermediate Representation" Compiler Infrastructure I see that this is the mother-project of the current MLIR which is included with LLVM (and that it was moved). But I would assume the code is no longer maintained, and thus would not compile with the current MLIR code.

Please, help. What I want is to:

  1. Have standard definitions and lowerings of the basic TF operations (such as Conv2D).
  2. Possibly be able to parse TF code into MLIR.


This does exist in parts across several projects and with no adequate documentation. And as you noted, it appears that a naive search on Google reveals the old, moved repo. The correct top level directory in the tensorflow repo is: tensorflow/tensorflow/compiler/mlir at master · tensorflow/tensorflow · GitHub

I know that @bondhugula has a flow that works, and there was recently a long thread on IREE’s discord where we helped a motivated user extract what we had done. In our last meetup, he discussed pulling some kind of documentation together with findings and approach, but I am not aware of anything publicly yet (to a level of completeness that will be helpful).

The above (correct) directory does reveal a number of starting points, and I suspect that many people here can help with specific questions. The two tools you are looking for to start out with are tf-translate to convert in/out of MLIR and tf-opt.

(Disclaimer: I am primarily a user of tensorflow, not involved with what is prioritized in this area)

1 Like

Yes, see tensorflow/tensorflow/compiler/mlir/tensorflow at master · tensorflow/tensorflow · GitHub (dialect in IR directory) and 'tf' Dialect  |  TensorFlow MLIR for the generated docs.

The primary lowerings is towards MHLO and then from there to LinAlg, SCF, Affine. But different approaches across here and in IREE too I believe. As Stella mentioned we need to pull the documentation for this together still. @mehdi_amini also still has a pending date to talk about this at ODM.

See tensorflow/import_model.h at master · tensorflow/tensorflow · GitHub and the tf-mlir-translate tool (you can see its usage for both GraphDef and SavedModel in the test directory).

You can also ask questions on mlir@tensorflow.org

1 Like

Your assumption isn’t right! MLIR is used via the llvm-project external repository that TensorFlow depends on. This is usually just a few weeks behind the llvm trunk. All MLIR/TF stuff is actively maintained/tested.

As @stellaraccident mentioned, please see under tensorflow/compiler/mlir. If you have a TF graphdef to try out for example, please see the thread here for information: https://groups.google.com/a/tensorflow.org/g/mlir/c/quu_s4KMzuM/m/Yn1ztQqqBAAJ
As far as the MLIR TF dialects go, you may want to typically post subsequent questions on the MLIR/TF forum.

Thanks @bondhugula @jpienaar @stellaraccident .
The code looks good.

What I cannot find is the compilation instructions for the code. I have installed bazel, but when I execute ./configure it never asks me where to search for LLVM/MLIR, and of course I have no idea how to launch the compilation (there is no top-level Makefile). Is there some quick guide on this somewhere? Or maybe you have a script that automates it?

Thank you in advance,

1 Like


bazel build --copt=-UNDEBUG --linkopt='-fuse-ld=lld'  tensorflow/compiler/mlir:tf-opt

But this is creating a completely new instance of LLVM/MLIR. Is there no way of making it depend on my pre-existing MLIR installation?

It gets updated to head about 2x a day during work week :slight_smile:

There is, but none documented there specifically: you’d have to change the bazel workspace file to point to local checkout instead, or you could try overriding the repo ( External dependencies - Bazel main) and would need to copy over the build files needed (the build files are not autogenerated nor is cmake used there; I haven’t used this approach though but could be least changes).

I normally just modify the checked out version from its cached location for small changes. Given the update frequency this is normally sufficient for me (rsync or just plain co from one to another). But I also mostly use a different build system and so haven’t tried much using bazel.

This would be a good addition to a MD doc there if folks find a good way that works well.

--override_repository=llvm-project=<your llvm-project copy>
(you’ll have to copy three BUILD files - one for llvm-project and two for mlir into your llvm-project repo - see tensorflow/workspace.bzl under llvm-project external dep).

1 Like

Added this as instructions in the README