Support XLA for both torch and tensorflow

[diegofiori]

Description

Currently nebullvm does not support TF-built in compiler XLA, which also allows the model to be compiled on Google TPUs. XLA is available for JAX, TF and PyTorch. Details about the TensorFlow and PyTorch implementations can be found in the following links:

• XLA for PyTorch
• XLA for Tensorflow

Integration

We should implement two new classes in nebullvm both inheriting from Compiler. They should implement the interfaces for PyTorch and TensorFlow of the XLA compiler. Note that XLA also supports reduced precision, so the quantization types must be managed too. For inspiration while coding the class interface you can analyse the Compilers built on top of TensorRT.

Note that XLA performance should be tested on a TPU for accessing its real performance.

TODO list

• [ ] Implement a PoC of the feature on Colab
• [ ] If not available a TPU, ping @valeriosofi for running the final version of the Colab notebook on Nebuly’s GCP account (Note that the debugging can be performed by yourself using other HW devices).
• [ ] Analyse the impact of the feature respect the actual AutoBoost
• [ ] If a positive impact of the feature is assessed we should implement it as a Compiler object else send your results to [email protected] detailing your findings and explaining why we should not proceed with the implementation of the feature
• [ ] Fork the nebullvm repo https://github.com/nebuly-ai/nebullvm
• [ ] Read the Contribution Guidelines
• [ ] Create a PR to main explaining your changes and showing the improvements obtained with the XLA compilers respect the previous versions

[comidan]

I'm gonna do it! 😁

[valeriosofi]

Hi @comidan, any news about this feature implementation? :)