[RFC] TF and TFLite legalizations to TOSA MLIR Dialect

[Suraj Sudhir]

TOSA Legalization from TensorFlow and TensorFlow Lite

Overview

This RFC describes the legalization passes that transform TensorFlow and TensorFlow Lite content to the TOSA dialect described in https://llvm.discourse.group/t/rfc-tosa-dialect-in-mlir/1971/18 . 

These passes implement the generation of validated sequences of TOSA operators corresponding to the input TensorFlow or TensorFlow Lite content. The legalizations express the functional sequence of TOSA operators required, as well as handling of numerical behavior, such as scaling, precision and saturation in the case of quantized integer datatypes.

Motivation

 The legalization passes from TensorFlow and TensorFlow Lite enables full networks to be legalized to a single, standard TOSA form. Both fp32 and quantized 8-bit networks have been exercised. A standard form of multiple high-level frameworks enables further code generation to be decoupled from the evolution of the high-level framework, and also insulates codegen from the differences between the operator specifications in different high-level frameworks.  

 The TOSA reference model described in the TOSA RFC enables TOSA output to be validated for bit accuracy against TensorFlow/TensorFlow Lite output. 

Components

 This implementation includes two pieces:

*   The actual MLIR lowerings from TensorFlow and TensorFlow Lite to TOSA, implemented in the form of multiple MLIR passes that may be registered with any MLIR pass manager.

*   Legalization test suites validating legalizations for TensorFlow and TensorFlow Lite on a per-operator basis for each implemented legalization from a high-level framework operator to one or more TOSA operators.

Implementation Details and Further Scope

 Approximately 100 TensorFlow and TensorFlow Lite operators are currently lowered to TOSA dialect form. These were chosen from occurrence frequency data generated from dozens of real-world networks in both TensorFlow and TensorFlow Lite , involving both fp32 and quantized 8-bit content.

The current set of lowerings enable the bit accurate lowering of nearly 50 real-world networks across several domains – image processing, object detection, super-resolution, speech recognition, speech synthesis, keyword spotting, natural language processing, and more.

Further legalizations may be added as needed, for operators that are not currently lowered to TOSA. The legalizations have not yet been exercised with networks that use quantized 16-bit int, fp16 or bfloat-16 datatypes. Most existing FP legalizations should work without alteration.  

Legalizations supporting the TOSA training profile are in development / under discussion. However tf.FakeQuantWithMinMaxArgs/Vals and tfl.quantize/dequantize op lowerings are currently present. 

Directory Setup

 The legalization code for both TensorFlow and TensorFlow Lite content are intended to sit within the TensorFlow repository at tensorflow/compiler/mlir/tosa .

Build System and Dependencies

The third_party/mlir/BUILD file within the TensorFlow repository will be updated to add a build object TosaOps that constructs the TOSA MLIR dialect library from the associated llvm-project repository. This is a build dependency for the legalization passes. The build dependencies of the TOSA dialect are described in the dialect RFC at https://llvm.discourse.group/t/rfc-tosa-dialect-in-mlir/1971/18 . 

TensorFlow and TensorFlow Lite dialects are build dependencies of the TOSA legalization passes. No other dialect currently present in the TensorFlow repository depends upon TOSA.  

[Stella Laurenzo]

As discussed offline around the review of the upstream MLIR RFC, I am strongly supportive of landing this part in the TensorFlow repo. Aside from being directly usable, having it in-repo will give us the tools over time to move code back and forth between the other TF/TFLite conversion sources and sinks to the benefit of all -- while also giving us a well-specified reference suite and alignment with code generation systems.

When it comes down to the actual PR and dependencies, you'll probably need the help of someone on the Google side, especially as our relationship between TF/LLVM has some "hidden nodes" that need some extra privileges to resolve. Me/the IREE team can help with that for the initial contribution if needed.

--
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/71198d33-efe3-4304-81f3-0ff6393996acn%40tensorflow.org.

[Suraj Sudhir]

Thank you, Stella! We'd gladly take any guidance and assistance with the process that we can get, as it's new to us and we want to avoid stepping on anyone's feet or breaking something. 

With all the other pieces out now, we're busy getting through the code release process and will coordinate with you when it's ready to go. 

[Jing Pu]

Thanks for the RFC! It looks like awesome work. I think my team at Google (Edge TPU) will be interested in making use of it.

I wonder if there is going to be an TOSA compiler available that can compile the "nearly 50 real-world networks" say on ARM CPU. 

> Legalization test suites validating legalizations for TensorFlow and TensorFlow Lite on a per-operator basis for each implemented legalization from a high-level framework operator to one or more TOSA operators.

This is also an interesting work in isolation. We are thinking about doing something similar and thinking of adding a general interpreter framework in MLIR to facilitate it, but haven't actually gone far enough to have a proposal for it. I will be interested in looking at how you approach it.

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/eda609bc-f194-4d07-9f5d-bd0aff76412bn%40tensorflow.org.

[Suraj Sudhir]

Thank you for your interest! E2e compilation - supporting different input frameworks and codegen paths constructed below it - is the goal that TOSA attempts to support. 

The networks we ran were in the context of generating TOSA from TF/TFLite input (fp32 as well as quantized int) - with output in TOSA form that matches bit exactly with the output from regular Tensorflow/TFLite. This is part of the legalization validation test setup, to ensure the TOSA form maintains fidelity. 

We are really excited that the Edge TPU team is interested in our work, and hope to continue this interaction. 

Suraj

[Suraj Sudhir]

The TOSA dialect implementation landed in LLVM last weekend. We are working on the remaining steps to open sourcing these legalization pieces now. 

[Jacques Pienaar]

Hey,

This sounds very promising. And I think it would fit well. As others have mentioned there are some questions wrt MHLO (especially it appears that there would be some savings to be had given the existing TF to MHLO legalizations), dynamic shapes and reuse in general. But that is something that could be worked on together. I was thinking we could even consider using either separate builds or configurable attributes (https://docs.bazel.build/versions/master/configurable-attributes.html) to make this easier initial integration/avoid it being on critical path from the start and so make it easier to start with.

Best,

Jacques

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/mlir/cd973ce8-752a-45df-a455-b9b783c29849n%40tensorflow.org.

[Suraj Sudhir]

(posting again - it seems the previous response only went to iree-discuss) 

Sure, we are happy to take feedback and guidance on how the build rules should look like here. We've previously used Bazel configurable attributes internally so we have at least some prior experience with using them . We will attempt to make the TOSA legalizations configurable, and then fine tune this with your inputs during the review process. 

[Suraj Sudhir]

We have the PR for this now: https://github.com/tensorflow/tensorflow/pull/44851

This version doesn't have any conditional build options set up. After talking to Stella Laurenzo, we were encouraged to create a PR of this content to begin with. 

[Suraj Sudhir]

With significant support from Stella who did detailed code review and helped get us thru TensorFlow contribution process concerns, this work has landed. 

Almost all of the original work has been contributed. The control flow legalization passes were deferred to another push, and will follow once review feedback is incorporated.