[C++][XLA Binding] executable serialization + allocator/compiled stats by csvance · Pull Request #3277 · EnzymeAD/Reactant.jl

csvance · 2026-09-12T18:29:43Z

This PR exposes executable serialization for PJRT so programs can be re-used without needing any kind of compile. This matters for ReactantServer.jl because every time the server restarts it can mean compiling 100+ mlir bundles, and even with an auto-tune cache this can take 30+ minutes.

It also exposes more detailed allocation usage and allows for clearing the memory high water mark for applications which need to measure exact memory water marks such as ReactantServer.

Without these provided by Reactant, I have to maintain my own downstream PJRT C bindings or patched Reactant jll which obviously isn't ideal in the long run. I also included IFRT bindings for the same functionality rather than just raising some sort of not implemented exception.

Working on building and testing this on my end. I confirmed that the general approach I took here works inside of ReactantServer via my vendored PJRT interface but it's one thing to just do all of the PJRT calls on my end and actually get this successfully upstreamed. Currently building Reactant jll locally so I can do a full end-to-end test of the changes.

csvance · 2026-09-12T19:22:47Z

I found one problem in the IFRT bindings with my local build, will push a fix shortly. I'm currently testing these changes inside ReactantServer with a huge model repository.

csvance · 2026-09-12T20:43:46Z

I have a ReactantServer feature branch which uses the the API changes here successfully, Server startup time is down from nearly 30 minutes to just under 3: EnzymeAD/ReactantServer.jl@main...feature/reactant-executable-cache

Not sure about CI / whether I can get the tests to pass before the new Reactant jll is built. I have my own internal Bazel build pulling the patched Reactant jll version here: https://github.com/csvance/Reactant.jl/releases/tag/reactantextra-local-06c35e0d

wsmoses · 2026-09-16T01:18:07Z

                             compile_options_proto_size));
 }

+#pragma region PjRtLoadedExecutable serialization and memory stats


what is pragma region

@wsmoses it looks like it was used many times elsewhere in API.cpp, probably because the file is so large and it serves as a way to collapse the block in an IDE? I can remove it if it would make the changes sit cleaner in the file or try and group things differently.

csvance · 2026-09-16T11:49:34Z

@wsmoses the tests I add pass with the local libReactant / Build Reactant_jll and I'm actively using the changes as part of a test deployment, everything is working as expected. There are some pre-existing test failures that show up in two of the local libReactant test builds, but those were there before my changes.

Merging this in would open up the potential to have a better kernel cache in Reactant beyond just autotuning. It's possible that I could eventually move ReactantServer off of the internal jll PJRT API and onto the public Reactant interface. These changes would be the first step towards that.

Pangoraw · 2026-09-16T12:19:06Z

Curious if you tried with using the compiled artifact with multi-gpu ? There was some problems last time i tried in #2779.

csvance · 2026-09-16T12:31:28Z

Curious if you tried with using the compiled artifact with multi-gpu ? There was some problems last time i tried in #2779.

I will test it out shortly; should be able to rule out whether there is some sort of issue on the XLA side since that's the only thing my PR touches.

csvance · 2026-09-16T15:31:13Z

@Pangoraw I tried a few different variations with both PJRT and IFRT on multiple GPUs (nVidia A6000):

Serialize on GPU 0,1, load on GPU 0,1
Serialize on GPU 0,1, load on GPU 1,0 (change order w/ CUDA_VISIBLE_DEVICES)

I don't currently have a GPU free to test something like 0,1 -> 1,2 or 0,1 -> 2,3.

The output of the compiled program was the same no matter which order it was loaded on. I only tried on a single machine; I don't have a cluster where I could try something distributed.

csvance · 2026-09-17T20:38:51Z

@wsmoses should I go ahead and merge this? Unsure of the process since it involves a bump to the jll.

Pangoraw · 2026-09-18T11:33:21Z

should I go ahead and merge this? Unsure of the process since it involves a bump to the jll.

to trigger a new jll build, you need to make a PR to https://github.com/EnzymeAD/ReactantBuilder targeting the relevant commit on Reactant main and bumping the Reactant_jll version.

Pangoraw · 2026-09-21T07:36:51Z

+    bytes = len == 0 ? UInt8[] : copy(unsafe_wrap(Array, ptr, (Int(len),); own=false))
+    @ccall free(ptr::Ptr{UInt8})::Cvoid


why not ?

Suggested change

bytes = len == 0 ? UInt8[] : copy(unsafe_wrap(Array, ptr, (Int(len),); own=false))

@ccall free(ptr::Ptr{UInt8})::Cvoid

return unsafe_wrap(Array, ptr, (Int(len),); own=true)

@Pangoraw I'll extract it from the helper in Utils.jl and just apply the unsafe_wrap(...; own=true) pattern to the two call sites. unsafe_wrap already correctly handles len == 0 here, better to just simplify.

wsmoses

also the api.cpp changes and the julia changes should be separate prs

csvance · 2026-09-21T13:57:47Z

@wsmoses I made this PR the C++ side only and extracted the Julia side into #3305

The only changes were running clang tidy so the code-style-cpp check passes.

…mpiled memory stats XLA can serialize a compiled executable and load it again later, reset a device allocator's peak counters, and report the compiler's static memory accounting for an executable, but none of that was reachable through libReactantExtra. This adds the C++ side only; the Julia bindings, API and tests follow in a separate PR once a Reactant_jll has been built from this. - PjRtLoadedExecutableSerialize / PjRtClientLoadSerializedExecutable (optional CompileOptionsProto override so a program can be placed on another device) - PjRtDeviceClearMemoryStats - PjRtLoadedExecutableGetCompiledMemoryStats with a flat JLCompiledMemoryStats struct - IFRT counterparts for the PJRT-backed IFRT client, going through IFRT's own LoadedExecutable::Serialize and Compiler::DeserializeLoadedExecutable - GetCpuPjrtApi, linking @xla//xla/pjrt/c:pjrt_c_api_cpu_internal, so the CPU plugin's PJRT C API table is reachable next to the GPU one - Darwin exported-symbol entries for the new names Built locally with deps/build_local.jl (CUDA) and exercised through the Julia bindings on CPU, on two RTX A6000s, and under the IFRT runtime. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0187bPmjLReauAiN6YP2a3cd

…ory stats (Julia) Julia side of the shims added to libReactantExtra in EnzymeAD#3277: XLA.serialize_executable, XLA.load_serialized_executable, XLA.clear_memory_stats!, XLA.compiled_memory_stats and the CompiledMemoryStats struct, with PJRT and IFRT methods and Thunk conveniences for serialize and compiled stats. The serialized bytes take ownership of the malloc'd buffer through unsafe_wrap(...; own=true) instead of being copied. test/core/executable_serialization.jl: round trip with bit-identical results, compile-options override, invalid bytes, memory stats, stats reset. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0187bPmjLReauAiN6YP2a3cd

csvance mentioned this pull request Sep 12, 2026

Persistent Kernel Cache EnzymeAD/ReactantServer.jl#14

Open

csvance force-pushed the feature/executable-serialization-memory-stats branch from 2447a46 to 06c35e0 Compare September 12, 2026 20:17

csvance marked this pull request as ready for review September 12, 2026 20:44

csvance requested a review from wsmoses September 16, 2026 01:14

wsmoses reviewed Sep 16, 2026

View reviewed changes

wsmoses approved these changes Sep 17, 2026

View reviewed changes

Pangoraw reviewed Sep 21, 2026

View reviewed changes

wsmoses reviewed Sep 21, 2026

View reviewed changes

csvance force-pushed the feature/executable-serialization-memory-stats branch from 0d63a6a to 3238dd5 Compare September 21, 2026 13:50

csvance changed the title ~~[XLA Binding] executable serialization + allocator/compiled stats~~ [C++][XLA Binding] executable serialization + allocator/compiled stats Sep 21, 2026

csvance mentioned this pull request Sep 21, 2026

[JL][XLA Binding] executable serialization + allocator/compiled stats #3305

Open

wsmoses force-pushed the feature/executable-serialization-memory-stats branch from 3238dd5 to 940d651 Compare September 21, 2026 17:00

wsmoses merged commit d420ec7 into EnzymeAD:main Sep 21, 2026
3 of 8 checks passed

		bytes = len == 0 ? UInt8[] : copy(unsafe_wrap(Array, ptr, (Int(len),); own=false))
		@ccall free(ptr::Ptr{UInt8})::Cvoid

	bytes = len == 0 ? UInt8[] : copy(unsafe_wrap(Array, ptr, (Int(len),); own=false))
	@ccall free(ptr::Ptr{UInt8})::Cvoid
	return unsafe_wrap(Array, ptr, (Int(len),); own=true)

Conversation

csvance commented Sep 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

csvance commented Sep 12, 2026

Uh oh!

csvance commented Sep 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wsmoses Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

csvance Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

csvance commented Sep 16, 2026

Uh oh!

Pangoraw commented Sep 16, 2026

Uh oh!

csvance commented Sep 16, 2026

Uh oh!

csvance commented Sep 16, 2026

Uh oh!

csvance commented Sep 17, 2026

Uh oh!

Pangoraw commented Sep 18, 2026

Uh oh!

Pangoraw Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

csvance Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

wsmoses left a comment

Choose a reason for hiding this comment

Uh oh!

csvance commented Sep 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

csvance commented Sep 12, 2026 •

edited

Loading

csvance commented Sep 12, 2026 •

edited

Loading