Segfault in XLA_CPU_Kernel Execute when calling jax compiled PJRT_LoadedExecutable_Execute

[duckmopsi]

Hello everyone,

I'm trying to execute a jax executable using PJRT C API (CPU Plugin) manually. In general, I'm following the steps given in this discussion and was able to replicate everything up until the execution step.

After extensive gdb debugging, I traced the crash to the following line from xla/xla/backends/cpu/runtime/kernel.h:

XLA_CPU_KernelError* error = (*kernel_)(&call_frame);

I assume I'm doing something wrong while handing over the execute_arguments but I'd really like to understand how XLA expects the arguments.

Error Context

The function I'm testing is:

def dummy(x, y):
    return jnp.dot(x, y) / 2.0

Input: two (2x2) float32 tensors
Output: one (2x2) float32 tensor

Therefore, I build two input buffers and fill them with dummy inputs ((1., 2., 3., 4.) and (2., 3., 4., 5.)). I also provide memory for the output PJRT-Buffer*.

I validated the args parameter at the final function call, as handed over in:

XLA_CPU_KernelCallFrame call_frame = {&num_workgroups, &workgroup_id, args.size(), args.data()};

• args.size() is 3
• the first 2 pointers of args.data() are the flattened input buffers (correct)
• the third pointer has the same dimension as the first two, but is filled with negative values (one value over all columns)
• the crash happens inside of the first thunk operation which is the .dot()-function

From my understanding, the third pointer is the output pointer. At this point I'm unsure whether:

1. my argument construction is wrong
2. the expected layout for CPU kernels is different, or
3. something internal is going wrong during the construction of args.

Environment

• Ubuntu-22.04 WSL2
• Python 3.11.13
• JAX Version 0.7.2
• XLA git repo on commit b28d5cae270a227b2db7daaceea20aff1248a904
• build command for .so file: bazel build //xla/pjrt/c:pjrt_api_c_cpu_plugin.so -c dbg inside of the xla root repository

How I generated the executable:

import jax
import jax.numpy as jnp
from jax import export
import functools
import os

def jax_to_ir(fn, input_shapes, constants=None, dest='model'):
    if not constants:
        constants = {}

    args = [jnp.zeros(s.shape, s.dtype) for _, s in input_shapes]
    fn_curried = functools.partial(fn, **constants)

    def ordered_wrapper(*args):
        arg_names = [arg_name for arg_name, _ in input_shapes]
        return fn_curried(**dict(zip(arg_names, args)))

    os.makedirs('export', exist_ok=True)
    
    comp = jax.jit(ordered_wrapper).lower(*args).compile()
    backend = jax.default_backend()
    xb = jax.lib.xla_bridge.get_backend(backend)

    serialized = xb.serialize_executable(comp.runtime_executable())
    with open('export/model_pjrt.bytes', 'wb') as f:
        f.write(serialized)

def dummy(x, y):
    return jnp.dot(x, y) / 2.0

jax_to_ir(dummy, [("x", jax.core.ShapedArray((2, 2), jnp.float32)), ("y", jax.core.ShapedArray((2, 2), jnp.float32))])

How I run it in C++:

After loading the plugin and api and creating an api client, I deserialize and load the executable as follows:

std::ifstream f("../python/export/model_pjrt.bytes", std::ios::binary | std::ios::ate);
std::streamsize sz = f.tellg();
f.seekg(0, std::ios::beg);

std::vector<char> buffer(sz);
if (!f.read(buffer.data(), sz)) {
    std::cerr << "Error reading file!" << std::endl;
    return 1;
}
f.close();

PJRT_Executable_DeserializeAndLoad_Args load_args;
memset(&load_args, 0, sizeof(load_args));
load_args.struct_size = PJRT_Executable_DeserializeAndLoad_Args_STRUCT_SIZE;
load_args.client = client;
load_args.serialized_executable = buffer.data();
load_args.serialized_executable_size = buffer.size();

error = api->PJRT_Executable_DeserializeAndLoad(&load_args);
PJRT_LoadedExecutable* exe = load_args.loaded_executable;

Then, after getting the execution device, I get the input buffers as follows, doing the same for my x and y input:

PJRT_Client_BufferFromHostBuffer_Args x_buf_args;
memset(&x_buf_args, 0, sizeof(x_buf_args));
x_buf_args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE;
x_buf_args.client = client;
x_buf_args.device = device;

// dummy input
std::vector<float> x = {1.f, 2.f, 3.f, 4.f};
std::vector<int64_t> x_dimensions = {2, 2};
x_buf_args.data = x.data();
x_buf_args.type = PJRT_Buffer_Type_F32;
x_buf_args.dims = x_dimensions.data();
x_buf_args.num_dims = x_dimensions.size();

PJRT_Error* buf_err = api->PJRT_Client_BufferFromHostBuffer(&x_buf_args);
if (buf_err) {
   std::cerr << "Buffer call failed for Tensor x." << std::endl;
   return 1;
}
PJRT_Buffer* x_buffer = x_buf_args.buffer;

Then, after repeating this for the y-buffer, I finally call the Execute-function:

PJRT_LoadedExecutable_Execute_Args exec_args;
memset(&exec_args, 0, sizeof(exec_args));
exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE;
exec_args.executable = exe;
exec_args.num_devices = 1;
exec_args.num_args = 2;

PJRT_Buffer* device_args[2] = { x_buffer, y_buffer};
PJRT_Buffer** all_devices[1] = { device_args };
exec_args.argument_lists = all_devices;

PJRT_ExecuteOptions exec_options;
memset(&exec_options, 0, sizeof(exec_options));
exec_options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE;
exec_args.options = &exec_options;

PJRT_Buffer* output_buffer_device[1];
output_buffer_device[0] = nullptr;
PJRT_Buffer** all_devices_out[1] = { output_buffer_device };
exec_args.output_lists = all_devices_out;
exec_args.execute_device = device;

PJRT_Error* exec_err = api->PJRT_LoadedExecutable_Execute(&exec_args);

At this point, the execution crashes inside of the .so-plugin and does not return any PJRT-Error, but instead:

AddressSanitizer:DEADLYSIGNAL
=================================================================
==4121==ERROR: AddressSanitizer: SEGV on unknown address 0x000000000010 (pc 0x75842c8da013 bp 0x7fff828d5ff0 sp 0x7fff828d5fe0 T0)
==4121==The signal is caused by a READ memory access.
==4121==Hint: address points to the zero page.
    #0 0x75842c8da013  (<unknown module>)

AddressSanitizer can not provide additional info.
SUMMARY: AddressSanitizer: SEGV (<unknown module>)
==4121==ABORTING

Main questions

• What is the expected layout of the args vector passed into CPU kernels?
• Should there be three pointers? If so, what is the meaning of the third one?
• Is my construction of argument_lists/output_lists incorrect?

Any insight would be greatly appreciated, as I've been digging into this with gdb for quite a while.

Thanks in advance! If you need any more information, I'll happily provide more context.