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

[ScottTodd]

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 tensorflow/tensorflow#83174
2025-01-17	#19683 merged, IREE nightly releases and source builds are no longer compatible with existing TOSA .mlir files

Upcoming

Date	Event
2025-??	Next TensorFlow stable release (2.19.0?)
2025-02-17	Planned IREE 3.2.0 stable release (#19641)
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

1. Direct users to older versions of these packages:

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

2. Watch for changes to the TOSA MLIR dialect in MLIR and integrate them into IREE as they arrive
3. 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
4. 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

[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'?

[ScottTodd]

TODO ahead of our next stable release:

• 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