Running Jax without Jax, skipping the Python layer using PJRT

By Joachim Sødequist, Danish e-infrastructure consortium (DeiC)

Building AI models means working within the AI ecosystem, often one builds on top of AI libraries with advanced features such as: automatic differentiation, graph kernel fusion, etc. These features make it easier to develop AI models while maintaining good performance during training. This typically means working in Python using one of a handful of optimised libraries. A majority of the compute-intensive workload is offloaded to highly optimised low-level computation kernels running on heterogeneous hardware. This is the ideal situation for model training, as the Python overhead is dwarfed by the time spent elsewhere.

After the dust has settled and we have a well-trained model, we now want to run inference on the model to speed up predictions from traditional C, C++, and Fortran HPC codes. This is where the different ecosystems clash, representing an inverse two-language problem.

The ordinary two-language problem refers to a duality: scientific codes are faster to prototype in slower programming languages, and then computationally heavy parts are often rewritten or offloaded to faster but more rigid programming languages. For example, in Python, the well-known linear algebra library Numpy is a thin wrapper that delegates work to optimised low-level routines like OpenBLAS. The ubiquitous machine learning libraries like PyTorch or Jax naturally follow the same pattern, although with quite a bit more logic in the Python interface. These frameworks naturally want to be able to prototype fast, just like scientific codes. However, this means that traditional HPC codes need to interface from the fast programming languages back into the slow programming languages. This is called embedding and can be done with libraries such as PyBind11 with C++ or CFFI with C/C++.

The software stack to go from C to Python and back to C can become quite involved. Especially when Python libraries often need to be containerised on large HPC systems such as LUMI. Thus, it would be much easier and nicer if we could tap directly into the low-level interfaces at inference time.

The low-level interface of Jax is PJRT (Pretty much Just another RunTime) that is the heart of the OpenXLA project. It takes as input either an Ahead-of-time compiled serialised executable and runs it, or an Intermediate Representation (IR) code to compile and then runs it. However, using PJRT is not very straightforward. Although Jax uses PJRT under the hood and there exist standardised PJRT API’s, one will find that all documentation is aimed at hardware vendors. It is, of course, important to have strong hardware support such that Jax can utilise a wide selection of accelerators. But using PJRT as an outsider to the Jax project boils down to following the recipe of a GitHub comment. In this Blog post, we will follow this recipe and illustrate how we can utilise the PJRT library directly in C with a small Hello-world example, similar to how Jax uses it. You can find an accompanying git repository with all the code here.

Using the PJRT C interface with compiled Jax kernels

 

 

First, we need to get the PJRT plugin for the specific hardware and Jax version. The hardware we have on hand is a AMD Radeon RX 9060 XT GPU for which we can find the corresponding pre-compiled PJRT library in PyPI under the name jax-rocm7-pjrt==0.6.0. After pip installing this, we find the required plugin as a single shared object file located at .venv/lib/python3.12/site-packages/jax_plugins/xla_rocm7/xla_rocm_plugin.so

Moreover, we also need the header file pjrt_c_api.h that describe the interface, this one can be found in the relevant ROCM XLA branch located at xla/xla/pjrt/c/pjrt_c_api.h.

In order to start running, we need to first create a client and establish the devices seen by the client. We use just the single GPU as the device in this example.

#include <stdio.h>
#include <stdlib.h>
#include <dlfcn.h>
#include "pjrt/pjrt_c_api.h"

#define NUM_DEVICES 1
#define DEVICE_ID 0

static const PJRT_Api* api;
static PJRT_Client* client;
static PJRT_Device* device;

void PJRT_Init() {
  // Load in the PJRT plugin
  void* handle = dlopen("pjrt/xla_rocm_plugin.so", RTLD_NOW | RTLD_LOCAL);
  const PJRT_Api* (*get_api)() = dlsym(handle, "GetPjrtApi");
  api = get_api();

  // Create PJRT client
  PJRT_Client_Create_Args create_args = {
    .struct_size = PJRT_Client_Create_Args_STRUCT_SIZE
  };
  api->PJRT_Client_Create(&create_args);
  client = create_args.client;

  // Get devices
  PJRT_Client_Devices_Args devices_args = {
    .struct_size = PJRT_Client_Devices_Args_STRUCT_SIZE,
    .client = client
  };
  api->PJRT_Client_Devices(&devices_args);
  device = devices_args.devices[DEVICE_ID];
} Here, the static const PJRT_Api* api will be the entry-point into all the functionality of the PJRT library, whereas client and device will be nicely tagged along for most calls. The next, and probably most important key ingredient is how to get the actual Jax kernel into PJRT. This requires some manual pre-processing in Python, which is shown in the next section. For now, we simply assume that we have a binary blob that somehow contains a nice pre-compiled serialized executable. We start by loading it like any other file and de-serialize it using the PJRT:

PJRT_LoadedExecutable* Load_Kernel(const char* filename) {
  FILE* f = fopen(filename, "rb");

  fseek(f, 0, SEEK_END);
  long size = ftell(f);
  fseek(f, 0, SEEK_SET);

  char* buffer = malloc(size);
  fread(buffer, 1, size, f);
  fclose(f);

  PJRT_Executable_DeserializeAndLoad_Args exe_args = {
    .struct_size = PJRT_Executable_DeserializeAndLoad_Args_STRUCT_SIZE,
    .client = client,
    .serialized_executable=buffer,
    .serialized_executable_size=size
  };
  api->PJRT_Executable_DeserializeAndLoad(&exe_args);
  free(buffer);
  return exe_args.loaded_executable;
}


All we need now is the standard scaffolding used in GPU programming in order to actually use the loaded_executable Jax kernel. First, we construct PJRT compatible input and output buffers on the host and the device side. For this example we just create some dummy C arrays and copy them to the empty device buffer: void Copy_HtD(float* input, size_t n, PJRT_Buffer** buffer){
  PJRT_Client_BufferFromHostBuffer_Args buffer_args = {
    .struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE,
    .client = client,
    .data = input,
    .type = PJRT_Buffer_Type_F32,
    .dims = (int64_t[]){n},
    .num_dims = 1,
    .device = device
  };
  api->PJRT_Client_BufferFromHostBuffer(&buffer_args);
  *buffer = buffer_args.buffer;
}

Where we see that the interface is flexible enough to handle multi-dimensional data as well, however, we just have one-dimensional test data. Additionally, in this example we're using PJRT_Client_BufferFromHostBuffer for illustrative purposes. In a real integrated production run, one might already have allocated buffers on the device from the HPC application code, in which case one can use PJRT_Client_CreateViewOfDeviceBuffer to create zero-copy buffers. Nevertheless, we can now execute the kernel on our input: 

void Execute(PJRT_LoadedExecutable* executable,
             PJRT_Buffer* buffer_x, PJRT_Buffer* buffer_y,
             PJRT_Buffer* output_d[]) {
  PJRT_Buffer* input_d[2] = { buffer_x, buffer_y };
  PJRT_Buffer* const* const argument_lists[NUM_DEVICES] = { input_d };
  PJRT_Buffer** output_lists[NUM_DEVICES] = { output_d };

  PJRT_ExecuteOptions options = {
    .struct_size = PJRT_ExecuteOptions_STRUCT_SIZE,
  };

  PJRT_LoadedExecutable_Execute_Args execute_args = {
    .struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE,
    .executable = executable,
    .execute_device = device,
    .argument_lists = argument_lists,
    .num_devices = NUM_DEVICES,
    .num_args = 2,
    .output_lists = output_lists,
    .options = &options
  };

  api->PJRT_LoadedExecutable_Execute(&execute_args);
}

Above, we first moved the input buffer pointers into a contiguous array, then, we wrapped up the input and output buffers into two dimensional arrays of size [num_devices, num_args] and [num_devices, num_outputs], respectively. The integers num_args=2 and num_output=1 represent a number of potentially high dimensional array buffers per device. The fact that we can send data to different devices in the num_devices axis provides a lot of flexibility for sharding data across multiple devices.

The API is not explicit about which parameters are mandatory or optional, unfortunately. For example, we have constructed an empty PJRT_ExecuteOptions struct, which seems a bit redundant, but it is in fact a required parameter. Additionally, there are also features in the API which are optionally implemented by the hardware vendors themselves. Thus, it is not always obvious which functionality exists. For example, execute_args can also return an PJRT_Event** device_complete_events struct for tracking when asynchronous operations in different devices have completed, however this does not seem to be required in this example. We only need to deal with the asynchronous events when we need to transfer the data back to the host.

void Copy_DtH(PJRT_Buffer *src, float *dst, size_t dst_size) {
  PJRT_Buffer_ToHostBuffer_Args to_host_args = {
    .struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE,
    .src = src,
    .dst = dst,
    .dst_size = dst_size
  };
  api->PJRT_Buffer_ToHostBuffer(&to_host_args);

  PJRT_Event_Await_Args ready = {
    .struct_size = PJRT_Event_Await_Args_STRUCT_SIZE,
    .event = to_host_args.event
  };
  api->PJRT_Event_Await(&ready);

  PJRT_Event_Destroy_Args destroy_ready_event = {
    .struct_size = PJRT_Event_Destroy_Args_STRUCT_SIZE,
    .event = to_host_args.event
  };
  api->PJRT_Event_Destroy(&destroy_ready_event);
}

The Device-to-Host transfer is straightforward, we instruct which device buffer must go into which host buffer. However, if we try to print the host buffer immediately after calling the PJRT_Buffer_ToHostBuffer routine we will find the host buffer contains uninitialized data. We must call await on the data transfer PJRT_Event to wait until the asynchronous background data transfer is completed. Here, we also free the PJRT_Event struct as we don’t need it anymore. As we are approaching the end of the example, we also need to destroy the rest of the device data buffers, the loaded executable and the client.

void Destroy_Args(PJRT_Buffer** output_d, PJRT_Buffer* buffer_x, PJRT_Buffer* buffer_y,
                  PJRT_LoadedExecutable* executable)
{
  PJRT_Buffer_Destroy_Args destroy_device_output = {
    .struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE,
    .buffer = output_d[DEVICE_ID]
  };
  api->PJRT_Buffer_Destroy(&destroy_device_output);

  PJRT_Buffer_Destroy_Args destroy_x = {
    .struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE,
    .buffer = buffer_x
  };
  api->PJRT_Buffer_Destroy(&destroy_x);

  PJRT_Buffer_Destroy_Args destroy_y = {
    .struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE,
    .buffer = buffer_y
  };
  api->PJRT_Buffer_Destroy(&destroy_y);

  PJRT_LoadedExecutable_Destroy_Args destroy_exec = {
    .struct_size = PJRT_LoadedExecutable_Destroy_Args_STRUCT_SIZE,
    .executable = executable
  };
  api->PJRT_LoadedExecutable_Destroy(&destroy_exec);

  PJRT_Client_Destroy_Args destroy_client = {
    .struct_size = PJRT_Client_Destroy_Args_STRUCT_SIZE,
    .client = client
  };
  api->PJRT_Client_Destroy(&destroy_client);

} All-in-all we have a C program main that looks like any ordinary GPU programming "Hello World" example, that is now using the PJRT API to execute a static size Jax kernel. As such, if this Jax kernel corresponds to a Machine Learning model, we can now use this to integrate model inference into the loop of any program that supports interfacing with C.

int main(int argc, char** argv) {
  // Step 1: Initialize API, Client and Device
  PJRT_Init();

  // Step 2: Load compiled binary
  PJRT_LoadedExecutable* executable = Load_Kernel("foo.bin");
  size_t num_outputs = 1;
  size_t size_out = 4, size_in = 4;

  // Step 3: Prepare input / output data buffers
  float input_x[] = {1.0f, 2.0f, 3.0f, 4.0f};
  float input_y[] = {5.5f, 6.6f, 7.7f, 8.8f};
  float output_h[size_out];

  PJRT_Buffer* buffer_x = NULL, *buffer_y = NULL, *output_d[num_outputs];
  Copy_HtD(input_x, size_in, &buffer_x);
  Copy_HtD(input_y, size_in, &buffer_y);

  // Step 4: Execute the computation
  Execute(executable, buffer_x, buffer_y, output_d);

  // Step 5: Retrieve results
  Copy_DtH(output_d[DEVICE_ID], output_h, sizeof(output_h));

  printf("Output:   [%.1f, %.1f, %.1f, %.1f]\n",
         output_h[0], output_h[1], output_h[2], output_h[3]);
  printf("Expected: [7.5, 10.6, 13.7, 16.8]\n");

  // Step 6: Cleanup
  Destroy_Args(output_d, buffer_x, buffer_y, executable);
  return 0;
} The advantages of this approach is that the software stack can be greatly simplified. We don't need a separate high level programming language for just the Machine Learning part. We don't even need the large machinery of OpenXLA and so we have minimized the number of external dependencies. All we need is any available C compiler. Getting rid of large dependencies have the additional advantage of decreasing any possible overhead. This performance gain, however, does come with reduced flexibility. There are many features in Jax that you might want at inference time such as dynamically Just-In-Time compiling the Jax kernel with different sized input during execution. One might be tempted to use PJRT_Client_Compile_Args to compile the StableHLO IR, however, this cannot change the input sizes dynamically since the size is baked into the IR.

Extracting Jax executable binary for PJRT

 

In previous section, we briefly skipped past how to actually extract the ahead-of-time compiled Jax kernel as a PJRT compatible binary. Luckily, Jax has all the functionality required. We first Ahead-Of-Time compile the kernel, this is a standard feature, and following the documentation one quickly arrives at:

import jax

def f(x, y):
    return 2 * x + y

x_shape = jax.ShapeDtypeStruct((4,), jax.numpy.float32)
y_shape = jax.ShapeDtypeStruct((4,), jax.numpy.float32)
compiled = jax.jit(f).trace(x_shape, y_shape).lower().compile()


Now we need to serialize this compiled Jax kernel, however, we cannot use the Jax export machinery. This is because Jax.export will lower the kernel into StableHLO IR in a flatbuffers binary along with a bunch of metadata. This metadata can be used by Jax to further transform the StableHLO IR before it is ready to pass it to PJRT. Therefore, we need to dig into Jax and extract the internal xla_bridge that really produces PJRT compatible binaries. 

from jax._src import xla_bridge as xb

executable = compiled._executable.xla_extension_executable()
backend = xb.get_backend('rocm')
serialized = backend.serialize_executable(executable)

with open('foo.bin', 'wb') as fd:
    fd.write(serialized)

We pass the raw executable inside the compiled object to the xla_bridge back-end, which serializes the executable. Finally, we can now pass this to PJRT as shown above, and run "Jax" without Jax by skipping the Python layer.