OpenXLA and PyTorch 2.0

[kevint324]

Hi Folks,

Is there a roadmap about how to integrate XLA into PyTorch 2.0 with the new released compile infrastructure.

Thanks

[shauheen]

Thank you for your question. PyTorch and XLA have had an ongoing integration for a few years which you can find at https://github.com/pytorch/xla

PyTorch/XLA is part of the PyTorch foundation and is being maintained by multiple partners who use XLA devices. PyTorch/XLA integration of TorchDynamo is ongoing and will be 2.0 at the same time as PyTorch. As XLA transitions to OpenXLA, PyTorch/XLA will also transition to use OpenXLA and produce StableHLO.

If you have any questions or feature requests please feel free to share them on the PyTorch/XLA repo.

[kevint324]

Hi Shauheen,

Thanks for the reply. Really appreciate. A few more question

According to the block diagram does XLA fits in the "Others" green block at codegen backends layer?
And is there any plan for Pytorch XLA GPU support? Here it says "Coming soon in a future release!".

Thank you.

Kevin

[shauheen]

Thank you Kevin for following up. The diagram may not be exact (at least not in short term) but generally you are correct XLA is another backend. I recommend this post for more detailed information on the status of the integration with TorchDynamo. As mentioned in the post PyTorch/XLA already works on GPU, however as you see in the documentation we have not been providing official support yet.

[kevint324]

Thank you Shauheen. These information are helpful.