Status of PjRt plugins

[joelberkeley]

Hi, what's the current status of PjRt plugin support with IREE? I looked at the pjrt directory but the README needs some level of updating.

I'm interested primarily in support for darwin ARM64 with Metal, and Linux AMD64 with CUDA

If there are instructions for building shared libraries for these PjRt plugins that would be even better.

Note I'm not using Python - I'm effectively using C

[ScottTodd]

Hi, what's the current status of PjRt plugin support with IREE? I looked at the pjrt directory but the README needs some level of updating.

The code you linked has been kept reasonably current, but it hasn't been put under much load for some time. @PragmaTwice has been sending patches to update the code, CI workflows, and documentation recently.

I'm interested primarily in support for darwin ARM64 with Metal, and Linux AMD64 with CUDA

Baseline support for those configurations shouldn't be too much work, but you may encounter some feature gaps off of the golden path (things like multi-device / sharding / tensor parallelism, support for uncommon operators, etc.).

If there are instructions for building shared libraries for these PjRt plugins that would be even better.

Note I'm not using Python - I'm effectively using C

I'm not sure about the shared library story for PJRT in general and IREE's implementation of the API. You might also be interested in

• Exporting from JAX using jax.export to get programs that the full IREE toolkit (not just PJRT) can work with
• The "core" IREE runtime C APIs provided in this repository
• The "shortfin" inference library and serving engine some of us are building out in https://github.com/nod-ai/shark-ai/tree/main/shortfin that is written in C/C++, uses the "core" IREE APIs itself, and include a "libshortfin" shared library

[joelberkeley]

OK thanks. Multi-device is ideally what I'm after, but I'm interested without it too.

My workflow has no access to JAX - I'm actually building something that plays the same role as JAX. I've been using the shared lib PJRT plugins provided by the XLA, and am looking at what IREE has to offer in that way.

[PragmaTwice]

Yeah, I think it's totally fine that you want to use the PjRt plugin without JAX. The plugin implementation itself (the built .so) doesn't contain any JAX-specific thing, for now.

You can follow the instruction in the README of the PjRt dir, i.e. pip install -v -e ... (or pip wheel --no-clean -v if you want) to build python packages, and then you can find a binary pjrt_plugin_iree_cpu.so in the build dir (if you choose python_packages/iree_cpu_plugin to build; linux assumed here). The entrypoint symbol GetPjrtApi is available in this shared library, which can be found by the linker or dlsym.

If you don't want to build the whole python package, but just the shared library. You can first configure the build via cmake, and then run the generated build script. e.g. (note that, in such case, maybe you need to build & configure the shared library of IREE compiler manually so that the PJRT plugin can find the compiler.)

# in the PJRT dir
mkdir build && cd build
cmake .. -DIREE_HAL_DRIVER_CUDA=ON -G Ninja # if you want to build the cuda plugin
ninja

(You can find additional CMake arguments in setup.py of each python package.)

Note that the cpu and cuda PJRT plugin is available, but the metal plugin is not. Although I think it should be easy to make a metal one (since it's supported by IREE compiler & runtime), I currently don't have any apple device in my hand so that it's a little hard for me to do this work.

[stellaraccident]

Speaking up as one of the original authors, this part of the project is currently "legit open source" -- ie. As a colleague of mine used to say: if it breaks, you own both pieces :)

My team at AMD, which maintains a lot of the project generally tries to help answer questions and help on this stuff, but afaik no one is currently paid to keep the pjrt plugins running.

With that said, I originally wrote it to be used and am not aware of any intrinsic reason why it shouldn't work. The contributions from folks like @PragmaTwice are appreciated, and we're happy to answer any questions, etc. And it's the kind of thing we may have more need to staff in the future. The maintenance/uplift work is appreciated even if the main support we can provide is moral.

[stellaraccident]

Hey @PragmaTwice, I just got your full handle. You've been showing up in my feed as just "Twice" for a while and I just realized your full handle is a play on "#pragma once". Well done and nice handle :)

[PragmaTwice]

Bingo! I started using "PragmaTwice" as my online ID when I began learning programming years ago. Over time, more people began to simply call me "Twice", so I eventually also accepted that as my nickname.