Third party XLA backend support

[Hengwen Tong]

❓ Questions and Help

Hi folks,

First, Happy new year! Hope you had a wonderful vacation.

We are working on a proprietary AI accelerator and we already have a working XLA backend within Tensorflow.
And next we are trying to enable our chip in PyTorch through XLA.

I know PyTorch/XLA supports TPU and GPU, and recently AWS announced their Trainium chip and it seems they too use XLA to hook in PyTorch.

I'm asking if there is any design documents to share to help us quickly understand the layers and blocks of PyTorch/XLA design.
I've read some design documents about the LTC . I've got a few questions:

1. From a vendor perspective which pieces of code should we look into to understand how to hook in a new accelerator though LTC?
2. Is there any extra work with the new announced Dynamo path?
3. As for the OpenXLA project, XLA is moving out of TF repo. So in some time PyTorch/XLA will depend on the OpenXLA repo instead of TF repo, is it? If yes, when might this happen? Q1? Q2?
4. Right now it seems the lazy tensor IR is converted to XLA HLO IR without MLIR on PyTorch side. The OpenXLA purposed a new StableHLO as a entry IR. Will PyTorch/XLA adopt StableHLO in some way?
5. From some docs I see the XRT the being deprecated. So is it a good idea that we just start with the PJRT?

Your reply would be greatly appreciated.

Thanks
Kevin