PJRT for Custom Accelarator

[Milin Bhade]

Hi — I’m working on a PJRT plugin for a custom accelerator and want to enable availability-aware, cost-driven partitioning of an XLA/HLO module across GPU, CPU, and the custom accelerator:

If only CPU + accelerator are available, run using those.

If GPU is present and used, automatically identify HLO subgraphs that are better offloaded to the accelerator and compile/run them there along with GPU

Questions:

Does XLA currently support multi-backend HLO partitioning/placement (i.e., splitting one HLO module across different backend types)?

Can a PJRT plugin expose device cost/constraints or otherwise influence partitioning during HLO-level compilation?

If not, what’s the recommended approach: implement an XLA pass to consume cost info, or build an orchestration layer that partitions the model and invokes multiple PJRT clients/executables? Which option is more realistic today?

I can prototype either an XLA/HLO pass or an external orchestrator — looking for pointers, existing examples, or caveats.

Thanks.

[notebookx userx]

tt-xla and intel Gaudi XLA plugin show how to use PJRT; besides the PRJT docs are pretty good themselves; I'd say walk with the interface before trying on all the bells and whistles for a final product

[Jose Fonseca]

One way to go about this would be to "daisy chain" PJRT plugins.  That is, one could have a PJRT plugin which loads a GPU/CPU/CUSTOM PJRT pluginS underneath, and then farms graphs or portion of graphs to each.  The advantage of this approach is that it would work with JAX and PyTorch/XLA without changing these projects or XLA itself.   However this intermediate PJRT plugin would need a good cost model of all underlying PJRT plugins performance (or build such model automatically based on experiments.)

Jose

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/fee447cc-5d51-4b12-84dc-5185be1150afn%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.