PJRT Plugins for other accelerators

[avik-pal]

• Tensortorrent: Enhancement: Support for Tenstorrent Hardware Acceleration via tt-mlir #1235 (done in Create TT client via plugin for TensTorrent devices #1860)
• Graphcore: They support JAX, unsure how to get the PJRT plugin??
• Via IREE (not the best approach but still nice to have for benchmarking) https://github.com/iree-org/iree/tree/main/integrations/pjrt/python_packages
  • CPU
  • CUDA
  • Vulkan
  • ROCm (we should ideally build the rocm support for XLA)
• Intel GPUs https://opensource.googleblog.com/2023/06/accelerate-jax-models-on-intel-gpus-via-pjrt.html
• Metal: We already have the code but waiting on upstream fixes jax-metal: Installation according to apples guide does not work. jax-ml/jax#26968
• AWS Trainium: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/jax/setup/jax-setup.html

[giordano]

Reiterating the question in #1235 (comment): how does this work? What component in the vendor stack I should look to use?

[avik-pal]

Essentially, check if they support Jax (that pretty much means they implement the right APIs we need).

For tt, this would be building the tt-xla repo. See https://github.com/tenstorrent/tt-xla/blob/dca15289be6a26dbae3012cbeea6549c59702710/examples/simple_regression.py#L19 for example, we need the pjrt plugin binary. We get that .so from an artifact, and then we can register that in reactant.

[avik-pal]

For metal & TPUs, we download the python whls and extract the relevant plugin binaries

[avik-pal]

For intel https://github.com/intel/intel-extension-for-openxla/blob/b1f1fab20a7e1993132489900344461938df2cc5/xla/BUILD#L2 might be enough?

[wsmoses]

for portability of intel it's better for us to distribute in a jll, but there are some pre-built: https://pypi.org/project/intel-extension-for-openxla/

I will note that for ROCM specifically that is a within XLA thing that we might as well enable a version in the reactant proper jll for, like CUDA

[avik-pal]

I haven't checked in too much detail but for custom calls we might need https://github.com/intel/intel-extension-for-openxla/blob/b1f1fab20a7e1993132489900344461938df2cc5/xla/python/__init__.py#L63-L75

[wsmoses]

also aws trainium: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/jax/setup/jax-setup.html

[avik-pal]

openxla/xla#27188 hermetic oneAPI build

[oschulz]

Vulkan support could be very interesting - maybe less for mainline ML, but Vulkan currently seems to be the only non-vendor-specific platform for accelerated raytracing, which has become very interesting for physics experiment that require fast photon-transport simulations (like IceCube, DUNE and also an increasing number of many smaller experiments). Vulkan.jl should be able do handle the ray-triangle-intersection problem already. But for other parts of such codes that need GPGPU, Reactant could maybe do that (via IREE)? Direct Julia to SPIR-V is not quite in reach yet, apparently.

[mofeing]

Direct Julia to SPIR-V is not quite in reach yet, apparently.

mmm crazy proposal but can we raise LLVM IR to MLIR GPU dialect which then we lower to SPIR-V?

[giordano]

Graphcore: They support JAX, unsure how to get the PJRT plugin??

Searching for "pjrt plugin" in the graphcore organisation doesn't give much. Alsom IPUs are currently listed as "legacy products" on their website, so I'm not holding my breath about that.