Calling pre-compiled JAX code from C++

[joaospinto]

Hi everyone,

My question is related to #5337, but I have a few extra requirements.

Suppose I want to deploy a JAX model, to be called from a C++ program.
As discussed in the link above, JAX can export an HLO file, which can then be compiled via xla::PjRtClient::Compile.

In many applications (e.g. real-time systems, memory-limited/restricted systems, etc), it would be desirable to offload this compilation process to an offline process, and simply include/link a library with the pre-compiled code.

Is this feasible? Can it already be done? Would greatly appreciate an example or some pointers. Thanks!

[joaospinto]

Let me provide a concrete example (but OK to depart from this in your answers).
Suppose I want to solve an optimization problem with a library like https://github.com/joaospinto/primal_dual_ilqr or https://github.com/google/trajax on a real-time system (e.g. some microcontroller in charge of controlling some dynamical system).
I'm not only trying to deploy some DNN, I need something a lot more generic.

[skye]

It's feasible and can be done, although it's not a smooth path. I recommend precompiling a program using JAX/Python, saving the compiled executable, then using the PJRT C API to deserialize, load, and execute that program.

Here are some pointers (sorry don't have time to write this all up atm):

• PJRT C API: https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h
• TPU impl: https://storage.googleapis.com/jax-releases/libtpu_releases.html
• GPU impl: https://pypi.org/project/jax-cuda12-pjrt/
• CPU impl info: How to use the PJRT C API? openxla/xla#7038 (comment)
• Using the API: How to use PJRT C API v.2 openxla/xla#13733 (comment)
• Compiling in jax: https://jax.readthedocs.io/en/latest/aot.html (we don't expose the raw PJRT executable serialization, I think you call .xla_executable or something on the compiled object, then serialize like this: https://github.com/google/jax/blob/071da567fb685708cc75528c2a6a1bcb06734275/jax/_src/compilation_cache.py#L227)
• Loading in C/++: https://github.com/openxla/xla/blob/ff218349ff7d307ef24bad2d99a73865586dfe17/xla/pjrt/c/pjrt_c_api.h#L1569-L1584
• Loading input buffers: https://github.com/openxla/xla/blob/ff218349ff7d307ef24bad2d99a73865586dfe17/xla/pjrt/c/pjrt_c_api.h#L733-L779
• Execute: https://github.com/openxla/xla/blob/ff218349ff7d307ef24bad2d99a73865586dfe17/xla/pjrt/c/pjrt_c_api.h#L1366-L1403
• Transfer buffers back to host: https://github.com/openxla/xla/blob/ff218349ff7d307ef24bad2d99a73865586dfe17/xla/pjrt/c/pjrt_c_api.h#L1694-L1714

Hope this helps! Feel free to ask more questions.

[joaospinto]

Thanks for the pointers @skye. Can the deserialization, loading, and execution happen without incurring any dynamic memory allocation or other OS interactions (e.g. launching processes, etc, at least when not swapping compute devices)?

[joaospinto]

@skye Here's what I'm trying to do (this is a minimal example, easy to reproduce locally): https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_aot.cpp

I'm currently running into
Bad StatusOr access: INTERNAL: TfrtCpuClient::DeserializeExecutable proto deserialization failed.

I think this is because C++ expects the serialized object to be a serialized ExecutableAndOptionsProto proto, but Python serializes only the first field of that proto, i.e. serialized_executable, and not the second field, i.e. compile_options. Is my understanding correct? How can I get a serialized ExecutableAndOptionsProto from Python directly?

[joaospinto]

I updated https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_aot.cpp to take in a serialized ExecutableAndOptionsProto proto (which I build in C++, in https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_hlo_to_serialized.cpp), but it runs into this strange error:

2024-07-03 11:45:21.312474: E external/xla/xla/status_macros.cc:56] INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_module.cc:551) proto.has_host_program_shape() No program shape found in the proto
*** Begin stack trace ***
        tsl::CurrentStackTrace()
        xla::status_macros::LogError(absl::lts_20230802::Status const&, char const*, int, int, bool)
        xla::status_macros::MakeError(char const*, int, absl::lts_20230802::StatusCode, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&, bool, int, bool)
        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::status_macros::MakeErrorStream::GetStatus() const
        xla::status_macros::MakeErrorStream::MakeErrorStreamWithOutput::operator absl::lts_20230802::StatusOr<std::__1::unique_ptr<xla::HloModule, std::__1::default_delete<xla::HloModule>>><std::__1::unique_ptr<xla::HloModule, std::__1::default_delete<xla::HloModule>>>()
        xla::HloModule::CreateFromProto(xla::HloModuleProto const&, xla::HloModuleConfig const&, bool)
        xla::HloModule::CreateFromProtoWithConfig(xla::HloModuleProtoWithConfig const&, bool)
        xla::cpu::(anonymous namespace)::CpuExecutableAotCompilationResult::FromString(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&)
        xla::cpu::CpuCompiler::LoadAotCompilationResult(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&)
        xla::TfrtCpuClient::DeserializeExecutable(std::__1::basic_string_view<char, std::__1::char_traits<char>>, std::__1::optional<xla::CompileOptions>)
        main
        start
*** End stack trace ***

libc++abi: terminating due to uncaught exception of type absl::lts_20230802::BadStatusOrAccess: Bad StatusOr access: INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_module.cc:551) proto.has_host_program_shape() No program shape found in the proto```

[skye]

You mean on the host? Unfortunately it does tons of dynamic memory allocation -- the C++ code is just not written for a constrained host environment. It doesn't launch processes, but I suspect doing host/device memory transfers involves syscalls.

[joaospinto]

OK, thanks for clarifying. I suppose the error above is more of an XLA issue, not a JAX one. Created an issue there instead: openxla/xla#15022

The only remaining JAX-related question I have (that is relevant for this conversation) is this: #22266

Will close this discussion for now, feel free to reply there.