Enable PJRT C API Client for Intel XPU

[jzhoulon]

Hi, developers
I am from XLA team in Intel, we are developing an XLA compiler with Intel GPU support, the compiler is based on PJRT C API interface, we want to leverage the benefits of multi front-end support on PjRt Client and want to make it works on PyTorch(currently it can work on JAX and partially on TensorFlow, hope it can work on PyTorch also:) ).
As you may know, we have Intel's device support in Framework(Pytorch, TensorFlow, JAX) with device type as XPU, so we also want to leverage this device type in PyTorch/XLA.
This PR follows the similar way as TPU:

export XPU_LIBRARY_PATH="xx/libxpu.so"
export PJRT_DEVICE=XPU
python torch_model.py

Any input/comment is appreciate:)

[shauheen]

Thanks @jzhoulon for submitting this PR. This is great to see.

[will-cromar]

Wow! I'm really happy to see more usage of the PJRT C API.

It looks like a bunch of unrelated formatting fixes in pjrt_computation_client.cc got picked up as well since our formatting check was missing it. I sent a fix in #4892 that should fix it, so please rebase after that merges.

For my own education, can you tell me more about your test plan for XPU on PyTorch/XLA going forward? Will you be running any automated tests to make sure this doesn't break, for example when we update our TensorFlow pin as in #4840?

Does XPU support any sort of multiprocessing or workload distribution that would require updates in xmp.spawn?

[yiqianglee]

@shauheen @will-cromar Thanks for your quick response, yes, we are actively working on PjRT based OpenXLA integration in Intel GPU, we have open-sourced our first version already, starting to support JAX front-end, and then support PyTorch-XLA (after this PR merged). OpenXLA on Intel GPU.

For test plan, that's a good question, the plan is keep testing stable PyTorch-XLA version (like release version) in our regular validation, and keep monitoring master to check if there are any breaking change (preparing for next release), we have similar experience to co-work between TensorFlow + Intel extension for TensorFlow.

[jzhoulon]

@will-cromar Thanks for the review! I have done the rebase and the clang-format change disappeared. Yes, XPU will support multiprocessing. We start from single process support in this PR and will have following PR to support multiprocessing.

[will-cromar]

LGTM