Skip to content

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

Open
csvance wants to merge 5 commits into
EnzymeAD:mainfrom
csvance:feature/executable-serialization-julia
Open

csvance wants to merge 5 commits into
EnzymeAD:mainfrom
csvance:feature/executable-serialization-julia

Conversation

@csvance

@csvance csvance commented Sep 21, 2026

Copy link
Copy Markdown
Collaborator

Julia part of the changes originally contained in #3277

…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
csvance force-pushed the feature/executable-serialization-julia branch from adea978 to 799af3b Compare September 22, 2026 13:30
@csvance
csvance marked this pull request as ready for review September 22, 2026 14:25
@wsmoses

wsmoses commented Sep 22, 2026

Copy link
Copy Markdown
Member

Error in testset executable serialization:
Error During Test at D:\a\Reactant.jl\Reactant.jl\test\core\executable_serialization.jl:21
  Got exception outside of a @test
  MethodError: no method matching unsafe_wrap(::Type{Reactant.XLA.IFRT.Array}, ::Ptr{UInt8}, ::Tuple{Int64}; own::Bool)
  
  Closest candidates are:
    unsafe_wrap(::Union{Type{Array}, Type{Array{T}}, Type{Array{T, N}}}, ::Ptr{T}, ::Tuple{Vararg{Int64, N}}; own) where {T, N}
     @ Base pointer.jl:85
    unsafe_wrap(::Union{Type{Array}, Type{Array{T}}, Type{Array{T, N}}}, ::Ptr{T}, ::Tuple{Vararg{var"#s124", N}} where var"#s124"<:Integer; own) where {T, N}
     @ Base pointer.jl:96
    unsafe_wrap(::Union{Type{Array}, Type{Array{T}}, Type{Vector{T}}}, ::Ptr{T}, ::Integer; own) where T
     @ Base pointer.jl:90

@csvance

csvance commented Sep 22, 2026

Copy link
Copy Markdown
Collaborator Author

@wsmoses looks like a name collision, fixing now.

@Pangoraw Pangoraw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

some comments about unnecessary @preserve, may not be caused by this PR but would be nice to clean at some point (also for all the device.device)

Comment thread src/xla/IFRT/LoadedExecutable.jl Outdated
Comment thread src/xla/IFRT/LoadedExecutable.jl Outdated
Comment thread src/xla/IFRT/LoadedExecutable.jl Outdated
Comment thread src/xla/PJRT/LoadedExecutable.jl Outdated
Comment thread src/xla/PJRT/LoadedExecutable.jl Outdated
@csvance

csvance commented Sep 25, 2026 •

Copy link
Copy Markdown
Collaborator Author

Looks like a memory stats test is failing on TPU, taking a look at that now: https://github.com/EnzymeAD/Reactant.jl/actions/runs/36016679351/job/107692793412?pr=3305

Test Failed at /__w/Reactant.jl/Reactant.jl/test/core/executable_serialization.jl:93
  Expression: XLA.clear_memory_stats!(device)
    Expected: Reactant.XLA.ReactantInternalError
  No exception thrown

EDIT: the test is broken for TPU, rebasing and pushing a fix shortly.

Only the CPU device lacks allocator statistics; TPU supports the reset,
matching the existing allocatorstats test.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants