Stateful FFI instantiate function does not have access to GPU

[johnviljoen]

Hi - I am wanting to build on this example.

Specifically I am wanting to pass .Ctx<ffi::PlatformStream<cudaStream_t>>() to the instantiate function, which works for the execution function. However I have not been able to find the "switch" - if one exists - to tell XLA/JAX to allow for this to happen. The error I get is clear:

jaxlib._jax.XlaRuntimeError: INVALID_ARGUMENT: [instantiate] Failed to decode all FFI handler operands (bad operands at: 0)
Diagnostics:
Failed to get platform stream: XLA FFI GPU context is not available

Does anyone know where I need to tell the system to allow the xla ffi gpu context to be available in the instantiate phase?

My cpp code is here:

static ffi::ErrorOr<std::unique_ptr<State>> StateInstantiate(
  cudaStream_t stream,
  const int64_t* offset_ptr,
  int64_t offset_size
) {
  print_1d_gpu_data<int32_t>(offset_ptr, offset_size);
  return std::make_unique<State>(42);
}

static ffi::Error StateMultiply(cudaStream_t stream, State* state,
                                ffi::BufferR0<ffi::S32> multiplier,
                                const int64_t* offset_ptr,
                                int64_t offset_size,
                                ffi::ResultBufferR0<ffi::S32> out) {
  printf("we are in the multiply function\n");

  print_1d_gpu_data<int32_t>(offset_ptr, offset_size);

  int32_t multiplier_host;
  cudaMemcpyAsync(&multiplier_host, multiplier.typed_data(), sizeof(int32_t),
                  cudaMemcpyDeviceToHost, stream);
  cudaStreamSynchronize(stream);

  int32_t result = (state->value * multiplier_host);

  cudaMemcpyAsync(out->typed_data(), &result, sizeof(int32_t),
                  cudaMemcpyHostToDevice, stream);
  cudaStreamSynchronize(stream);
  return ffi::Error::Success();
}

XLA_FFI_DEFINE_HANDLER(kStateInstantiate, StateInstantiate,
                       ffi::Ffi::BindInstantiate()
                      .Ctx<ffi::PlatformStream<cudaStream_t>>()
                      .Attr<ffi::Pointer<const int64_t>>("ptr")
                      .Attr<int64_t>("size")
);

XLA_FFI_DEFINE_HANDLER(
    kStateMultiply, StateMultiply,
    ffi::Ffi::Bind()
        .Ctx<ffi::PlatformStream<cudaStream_t>>()
        .Ctx<ffi::State<State>>()
        .Arg<ffi::BufferR0<ffi::S32>>()
        .Attr<ffi::Pointer<const int64_t>>("ptr")
        .Attr<int64_t>("size")
        .Ret<ffi::BufferR0<ffi::S32>>());

NB_MODULE(_gpu_example, m) {
  m.def("type_id",
        []() { return nb::capsule(reinterpret_cast<void*>(&State::id)); });
  m.def("handler", []() {
    nb::dict d;
    d["instantiate"] = nb::capsule(reinterpret_cast<void*>(kStateInstantiate));
    d["execute"] = nb::capsule(reinterpret_cast<void*>(kStateMultiply));
    return d;
  });
}

my full Python is here:

import jax
import jax.extend.core
from my_package._gpu_example import handler, type_id
import jax.numpy as jnp
from jax.interpreters import mlir, batching
import functools as ft

jax.ffi.register_ffi_target("state", handler(), platform="CUDA")
jax.ffi.register_ffi_type_id("state", type_id(), platform="CUDA")

state_p = jax.extend.core.Primitive("state")
state_p.def_effectful_abstract_eval = True

@state_p.def_impl
def state_impl(x, ptr, size):
    call = jax.ffi.ffi_call(
        "state",
        jax.ShapeDtypeStruct((), jnp.int32),
        has_side_effect=True # makes sure it is executed even when outputs are not used
    )
    return call(x, ptr=ptr, size=size)

solve_single_f64_low = mlir.lower_fun(state_impl, multiple_results=False)
mlir.register_lowering(state_p, solve_single_f64_low)

@state_p.def_abstract_eval
def state_aval(x_aval, ptr, size):
    return jax.core.ShapedArray((), jnp.int32)

@ft.partial(jax.jit, static_argnames=["ptr", "size"])
def state(x, ptr, size):
    return state_p.bind(x, ptr=ptr, size=size)

# The state is 42. Let's multiply it by 10.
multiplier = jnp.int32(10)
offset = jnp.array([1, 2], dtype=jnp.int32) # (1,2,3) # list(jnp.array(1, dtype=jnp.int32))
ptr = offset.unsafe_buffer_pointer()
size = offset.size
result = state(multiplier, ptr, size)

Also is there any documentation on the phases - I can of course see there is instantiation and execution, but I have seen references to prepare and initialize What is the difference between these phases, what rules do they abide by, is there any documentation?