iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.87k stars 626 forks source link

Restructure Python wheels following TF[Lite] importer changes #13483

Open ScottTodd opened 1 year ago

ScottTodd commented 1 year ago

Following the work in https://github.com/openxla/iree/issues/13037 and other issues/PRs, our TensorFlow and TensorFlow Lite importers are now pure Python packages/scripts, with no direct C++ dependencies on TensorFlow. As a result, our iree_tools_tf and iree_tools_tflite-*.whl Python wheels are now significantly simpler to build and smaller in size (~50MB before -> ~30KB after). We also no longer have a iree_tools_xla-*.whl package.

I'd like to take this chance to see if we can further simplify or restructure our released Python packages.

I was also looking at adding badges to our repo README like these: PyPI PyPI for those, it would be nice to have a smaller list (ideally just one package).

Some ideas:

A: Merge all iree_tools_* packages into iree_compiler

Include the Python sources for the TF/TFLite tools in the main package, but make them optional by adding checks around code like this: https://github.com/openxla/iree/blob/060616cea34d062434097db155c458edb814e235/integrations/tensorflow/python_projects/iree_tf/iree/tools/tf/scripts/iree_import_tf/__main__.py#L60-L72

try:
  import tensorflow as tf
except ImportError as e:
  # print error / instructions for installing

Pros:

Cons:

B: Merge just iree_tools_tf and iree_tools_tflite

These packages are small and similar enough that they could be merged. If they have the same requirements (tensorflow/tf-nightly), is there any harm in merging them?

stellaraccident commented 1 year ago

Option B, merging both into iree-tools-tf, seems the right thing. We can't introduce TF dependencies to the compiler.

hcindyl commented 1 year ago

Sorry just caught up with this since the IREE nightly release finally has this change.

We need to update https://openxla.github.io/iree/getting-started/tflite/#using-command-line-tools if it is assumed tensorflow python package is installed. The alternative is to add tensorflow in setup.py so it is included when the wheels are installed.

The iree-import-tflite command line behavior at iree-tools-tflite-20230509.514

Traceback (most recent call last):
  File ".../iree-import-tflite", line 5, in <module>
    from iree.tools.tflite.scripts.iree_import_tflite.__main__ import main
  File ".../iree/tools/tflite/scripts/iree_import_tflite/__main__.py", line 14, in <module>
    from tensorflow.python.pywrap_mlir import experimental_tflite_to_tosa_bytecode
ModuleNotFoundError: No module named 'tensorflow'

Please advise on the proper flow to set the wheels up.

ScottTodd commented 1 year ago

The current source of truth for what dependencies is required is https://github.com/openxla/iree/blob/main/integrations/tensorflow/test/requirements.txt (as of May 9 2023, that is tf-nightly==2.13.0.dev20230501). Once the Python APIs (experimental_tflite_to_tosa_bytecode and experimental_convert_saved_model_to_mlir) are stable / available in non-nightly TF releases, setup should stabilize.

hcindyl commented 1 year ago

The current source of truth for what dependencies is required is https://github.com/openxla/iree/blob/main/integrations/tensorflow/test/requirements.txt (as of May 9 2023, that is tf-nightly==2.13.0.dev20230501). Once the Python APIs (experimental_tflite_to_tosa_bytecode and experimental_convert_saved_model_to_mlir) are stable / available in non-nightly TF releases, setup should stabilize.

Understood. In that case, in the webdoc https://openxla.github.io/iree/getting-started/tflite/#prerequisites should also include tensorflow (or point to the requirments.txt). Is that correct?

ScottTodd commented 1 year ago

Yep. Maybe we should just include tf-nightly in the recommended pip install instructions.