Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Upstream the ONNX importer. #2636

Merged
merged 5 commits into from
Dec 13, 2023
Merged

Upstream the ONNX importer. #2636

merged 5 commits into from
Dec 13, 2023

Conversation

stellaraccident
Copy link
Collaborator

This is part 1 of 2, which will also include upstreaming the FX importer. I started with ONNX because it forces some project layout updates and is more self contained/easier as a first step.

Deviating somewhat from the RFCs on project layout, I made the following decisions:

  • Locating the onnx_importer.py into torch_mlir.extras as Maks already has opened up that namespace and it seemed to fit. Better to have fewer things at that level.
  • Setup the build so that the root project only contains MLIR Python and pure Python deps (like the importers), but this can be augmented with the projects/ adding more depending on which features are enabled.
  • The default build continues to build everything whereas in TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS=1 mode, it builds a torch-mlir-core wheel with the pure contents only.

onnx_importer.py and importer_smoke_test.py are almost verbatim copies from SHARK-Turbine. I made some minor local alterations to adapt to paths and generalize the way they interact with the outer project. I expect I can copy these back to Turbine verbatim from here. I also updated the license boilerplate (they have the same license but slightly different project norms for the headers) but retained the correct copyright.

Other updates:

  • Added the ONNX importer unit test (which also can generate test data) in lit, conditioned on the availability of the Python onnx package. In a followup once I know everything is stable, I'll add another env var that the CI can set to always enable this so we know conclusively if tests pass.
  • Moved the ONNX conversion readme to docs/.
  • Renamed CMake option TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS -> TORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS and inverted the sense. Made the JitIR importer and LTC options cmake_dependent_options for robustness.

@stellaraccident stellaraccident marked this pull request as ready for review December 13, 2023 01:37
@stellaraccident stellaraccident merged commit 74f7a0c into main Dec 13, 2023
5 checks passed
@stellaraccident stellaraccident deleted the upstream_onnx_importer branch December 13, 2023 03:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants