iree Support for LiteRT (TensorFlow Lite, .tflite) with TOSA 1.0

Background

The TOSA MLIR dialect is one of IREE's accepted input formats, commonly used to compile and then run LiteRT / TensorFlow Lite programs. The TOSA MLIR dialect is in the process of incrementally upgrading to v1.0: https://discourse.llvm.org/t/rfc-tosa-dialect-increment-to-v1-0/83708. This will introduce a new backwards compatibility baseline after a series of breaking changes to operator definitions.

[!WARNING] During this upgrade process, there will be several planned and unplanned breaking changes that will affect user and developer workflows.

Timelines

Already happened

Date	Event
2024-10-24	TensorFlow 2.18.0 stable release: https://pypi.org/project/tensorflow/#history
2025-01-08	IREE 3.1.0 stable release: https://github.com/iree-org/iree/releases/tag/v3.1.0
2025-01-08	TensorFlow removed the `ExperimentalTFLiteToTosaBytecode` function in https://github.com/tensorflow/tensorflow/pull/83174
2025-01-17	https://github.com/iree-org/iree/pull/19683 merged, IREE nightly releases and source builds are no longer compatible with existing TOSA .mlir files
2025-02-10	IREE 3.2.0 stable release: https://github.com/iree-org/iree/releases/tag/v3.2.0 (https://github.com/iree-org/iree/issues/19641)
2025-02-21	TensorFlow 2.19.0rc0 release

Upcoming

Date	Event
2025-04	Planned deadline for TOSA v1.0 changes to be rolled out

Support plan

Given the existing and planned events, compatibility is expected to be:

For tensorflow<=2.18.0, use iree-tools-tflite<=20250107.1133 and iree-base-compiler<=3.1.0
For iree-base-compiler>=3.2.0 (including current source builds), TOSA .mlir files are likely to be unsupported as the IREE compiler has started to include TOSA v1.0 breaking changes and TensorFlow does not generate the new format
For tf-nightly or newer tensorflow versions, the API that iree-tools-tflite relies on no longer exists

Our current plan is to

Direct users to older versions of these packages:

tensorflow<=2.18.0
iree-tools-tflite<=20250107.1133
iree-base-compiler<=3.1.0

Watch for changes to the TOSA MLIR dialect in MLIR and integrate them into IREE as they arrive
Watch for a new .tflite to TOSA .mlir API to be added to TensorFlow (or another project like TFLite / LiteRT / etc.) then switch the iree-import-tflite script in the iree-tools-tflite package to use it
Advertise support for TFLite/TOSA in IREE again after these changes have fully rolled out and have been tested

This could change with different engineering decisions (e.g. carrying more reverts in IREE). Feedback and patches are welcome.

References

TFLite framework support page: https://iree.dev/guides/ml-frameworks/tflite/
input_tosa compiler plugin source: compiler/plugins/input/TOSA
iree-tosa-input-transformation-pipeline source: compiler/plugins/input/TOSA/InputConversion/Passes.cpp
iree-tools-tflite Python project source: integrations/tensorflow/python_projects/iree_tflite
iree-import-tflite script source: integrations/tensorflow/python_projects/iree_tflite/iree/tools/tflite/scripts/iree_import_tflite/__main__.py
TensorFlow source code for TOSA support: tensorflow/compiler/mlir/tosa
MLIR source code for TOSA support

Jan 23 '25 00:01 ScottTodd

Sample error message when using a recent iree-compile and an existing TOSA .mlir generated by the (out of sync) code in TensorFlow:

 INFO     litert_models.utils:utils.py:90 Launching compile command:
  cd /home/runner/work/iree-test-suites/iree-test-suites/litert_models && iree-compile /home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1.mlirbc --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-target-cpu=host -o /home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1_cpu.vmfb
ERROR    litert_models.utils:utils.py:96 Compilation of '/home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1_cpu.vmfb' failed
ERROR    litert_models.utils:utils.py:97 iree-compile stdout:
ERROR    litert_models.utils:utils.py:98 
ERROR    litert_models.utils:utils.py:99 iree-compile stderr:
ERROR    litert_models.utils:utils.py:100 <unknown>:0: error: loc("MobilenetV1/MobilenetV1/Conv2d_0/Relu6"): 'tosa.conv2d' op requires attribute 'acc_type'
<unknown>:0: note: loc("MobilenetV1/MobilenetV1/Conv2d_0/Relu6"): see current operation: %56 = "tosa.conv2d"(%arg0, %55, %54) <{dilation = array<i64: 1, 1>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 2, 2>}> : (tensor<1x224x224x3xf32>, tensor<8x3x3x3xf32>, tensor<8xf32>) -> tensor<1x112x112x8xf32>
<unknown>:0: note: loc("MobilenetV1/MobilenetV1/Conv2d_0/Relu6"): in bytecode version 1 produced by: MLIR20.0.0gi

Sample error message when using tf-nightly with iree-import-tflite calling the API that was removed:

Traceback (most recent call last):
  File "/home/ben/.local/bin/iree-import-tflite", line 8, in <module>
    sys.exit(main())
  File "/home/ben/.local/lib/python3.10/site-packages/iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 54, in main
    tflite_to_tosa(
  File "/home/ben/.local/lib/python3.10/site-packages/iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 69, in tflite_to_tosa
    tflite_to_tosa_bytecode(
  File "/home/ben/.local/lib/python3.10/site-packages/tensorflow/python/compiler/mlir/mlir.py", line 200, in tflite_to_tosa_bytecode
    pywrap_mlir.experimental_tflite_to_tosa_bytecode(
  File "/home/ben/.local/lib/python3.10/site-packages/tensorflow/python/pywrap_mlir.py", line 123, in experimental_tflite_to_tosa_bytecode
    return ExperimentalTFLiteToTosaBytecode(
NameError: name 'ExperimentalTFLiteToTosaBytecode' is not defined. Did you mean: 'experimental_tflite_to_tosa_bytecode'?

Jan 23 '25 00:01 ScottTodd

TODO ahead of our next stable release:

[x] Add a warning about compatibility to the top of https://iree.dev/guides/ml-frameworks/tflite/
[ ] (If possible with current APIs) add version check warnings/errors to https://github.com/iree-org/iree/tree/main/compiler/plugins/input/TOSA
[ ] expand on version check warnings/errors in https://github.com/iree-org/iree/blob/main/integrations/tensorflow/python_projects/iree_tflite/iree/tools/tflite/scripts/iree_import_tflite/main.py

Jan 27 '25 23:01 ScottTodd

TensorFlow 2.19.0 started rolling out on 2025-02-21: https://pypi.org/project/tensorflow/#history

This resulted in test failures since the ExperimentalTFLiteToTosaBytecode function no longer exists: https://github.com/iree-org/iree-test-suites/actions/runs/13474005120/job/37650922309

litert_models/tests/kaggle/tensorflow/mobilenet_v1_test.py::test_mobilenet_v1_0_25_224 
-------------------------------- live log call ---------------------------------
INFO     litert_models.conftest:conftest.py:17 model_path: /home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1.tflite
INFO     litert_models.utils:utils.py:58 Importing '/home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1.tflite' to '/home/runner/.cache/kagglehub/models/tensorflow/mobilenet-v1/tfLite/0-25-224/1/1.mlirbc'
ERROR    litert_models.utils:utils.py:67 Import of '1.tflite' failed!
ERROR    litert_models.utils:utils.py:68 iree-import-tflite stdout:
ERROR    litert_models.utils:utils.py:69 
ERROR    litert_models.utils:utils.py:70 iree-import-tflite stderr:
...
File "/home/runner/work/iree-test-suites/iree-test-suites/.venv/bin/iree-import-tflite", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/runner/work/iree-test-suites/iree-test-suites/.venv/lib/python3.11/site-packages/iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 67, in main
    tflite_to_tosa(
  File "/home/runner/work/iree-test-suites/iree-test-suites/.venv/lib/python3.11/site-packages/iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 82, in tflite_to_tosa
    tflite_to_tosa_bytecode(
  File "/home/runner/work/iree-test-suites/iree-test-suites/.venv/lib/python3.11/site-packages/tensorflow/python/compiler/mlir/mlir.py", line 200, in tflite_to_tosa_bytecode
    pywrap_mlir.experimental_tflite_to_tosa_bytecode(
  File "/home/runner/work/iree-test-suites/iree-test-suites/.venv/lib/python3.11/site-packages/tensorflow/python/pywrap_mlir.py", line 123, in experimental_tflite_to_tosa_bytecode
    return ExperimentalTFLiteToTosaBytecode(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NameError: name 'ExperimentalTFLiteToTosaBytecode' is not defined. Did you mean: 'experimental_tflite_to_tosa_bytecode'?

FAILED                                                                   [ [50](https://github.com/iree-org/iree-test-suites/actions/runs/13474005120/job/37650922309#step:6:51)%]

We can add new error handling to iree-import-tflite (for as long we keep it around) and change the test suite to treat failures there as IreeImportTfLiteException or a new exception type.

Feb 24 '25 16:02 ScottTodd

Subscribing to the issue because this is a blocker for update to https://github.com/iree-gd/iree.gd/issues/105

Feb 26 '25 19:02 fire

Hi @ScottTodd Is there any update on this? Does the user still need to use the old version to compile LiteRT?

May 28 '25 05:05 ebraheemabbas

As far as I know, yes, the older versions are still required as a replacement for the old import API is not yet available. I've heard of some progress on Discord and on GitHub like here: https://github.com/tensorflow/tensorflow/pull/92055 . cc @sjarus

Jun 09 '25 17:06 ScottTodd

We should be able to offer an update next week - I'm afraid the owners are both on vacation this week @ScottTodd

Jun 10 '25 15:06 sjarus

Hi @sjarus maybe i missed it , had you published any update?

Jun 23 '25 08:06 ebraheemabbas

Sorry for the delay. We are currently working on a thin stand-alone Python wrapper around the TFLite -> TOSA code in upstream TensorFlow.

It is not finished yet but we hope to make it available publicly in the next few months.

Jun 23 '25 15:06 Tessil

Sorry for the delay. We are currently working on a thin stand-alone Python wrapper around the TFLite -> TOSA code in upstream TensorFlow.

It is not finished yet but we hope to make it available publicly in the next few months.

Thanks @Tessil ! Any chance we can make this available in a proof of concept form in July ? There's a lot of interest here...

Jun 23 '25 22:06 sjarus

We have to still check a few things internally but I hope we can have some form of beta that we can share next week, but nothing certain.

Jun 26 '25 16:06 Tessil

The tool is now available at https://gitlab.arm.com/tosa/tosa-converter-for-tflite

Note that the tool is still in early development and a formal release should come up in the next few weeks.

Jul 07 '25 15:07 Tessil