Add an opt-in Metal array storage preference - #405
Conversation
JACC.array(x; storage = :shared) on the Metal backend copies host data into unified SharedStorage memory instead of the default PrivateStorage, dropping the private-staging blit on Apple Silicon. storage = nothing (the default) preserves current behavior exactly on every backend; other backends raise a clear ArgumentError for a non-nothing value. Implements the offer in JuliaGPU#402. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
aurascoper
changed the title
|
Addressed the generic backend preference design and narrowed the follow-up to Metal allocation behavior. Core JACC now stores backend options in a backend-keyed map, and Full suite: 223 passed on Apple Silicon with Julia 1.12.6 and Metal.jl 1.10.0. Focused Could you take another look? |
|
@aurascoper thanks for the contribution! @PhilipFackler let me know if this is ready so I can run CI |
|
Test this please |
| end | ||
|
|
||
| function JACC._array(::MetalBackend, x::AbstractArray) | ||
| if _array_storage() == "shared" |
There was a problem hiding this comment.
This should probably be run just once during module load. Can we store this as a module global?
PhilipFackler
left a comment
PhilipFackler
left a comment
There was a problem hiding this comment.
I think this is good to go. Thanks @aurascoper !
|
@williamfgc looks like cousteau has an issue with the hip library. We can probably ignore that for this PR. |
Closes PhilipFackler's unresolved review thread on
ext/MetalExt/MetalExt.jl: _array_storage() was re-reading the backend
preference Dict and revalidating it on every JACC._array call, on the
same array-allocation hot path this PR already measures overhead on.
A plain module-load-once cache is not safe here: set_backend is
documented (and tested by array_storage_preference in
test/backend/metal.jl) to apply storage changes live within the same
Julia session, without a restart. That test calls set_backend three
times mid-session, including two invalid-value cases that must throw
on the very next JACC.array call.
Instead, _EXT_PREFS_GENERATION (Ref{Int}, core JACC) is bumped on
both real mutation sites of _EXT_PREFS (_set_extension_preferences,
unset_backend's empty!) — grepped exhaustively, no other write site
exists. MetalExt's _array_storage() caches (generation, value) and
only recomputes on a generation mismatch, so the hot path is a cheap
Int comparison instead of a Dict lookup plus revalidation, while the
live-reload contract the existing test locks in is unchanged.
Verified on Apple Silicon (Julia 1.12.6, Metal.jl 1.10.0):
array_storage_preference 16/16, full suite 223/223 — matches the
pre-fix baseline exactly, no regressions.
|
Addressed the module-global caching note on A plain load-once cache wasn't safe as-is: Used a generation counter instead: Verified on Apple Silicon (Julia 1.12.6, Metal.jl 1.10.0): |
3 participants
Summary
set_backend(...; kw...)options in a generic backend-keyed preference mapMetalExtconsume and validate thestorageoption it ownsJACC.array(x)MtlArrayidentity for implicit and explicit Metal callsWith Metal's default settings, host-array copies use private storage. Projects can opt into shared storage with
JACC.set_backend("Metal"; storage = :shared).Motivation
This implements the storage-selection contribution discussed in #402 and the generic preference design requested in review.
Metal's
MtlArray(x)construction path allocates and copies the source array (Metal.jlarray.jl:287-296), with private storage selected by default (array.jl:193-205). On Apple Silicon, shared storage can avoid the private-storage transfer path.Measurement
On an Apple M4 with Julia 1.12.6 and Metal.jl 1.10.0, constructing 10,000,000
Float32elements produced these medians across 30 samples per path::sharedwas 52% lower than:privateEach construction was followed by
Metal.synchronize(). Laptop idle state was uncontrolled. GPU bandwidth for compute-bound workloads remains unmeasured.Design
Core JACC persists backend options verbatim, and
MetalExtownsstorageinterpretation and validation, so the portableJACC.array(x)signature stays backend-independent. Other preference lifecycle operations continue to use the existing Preferences.jl path.The Metal specialization preserves the pre-PR identity behavior for existing
MtlArrayvalues and restores the explicitJACC.array(backend, x)host-copy method.MetalExt's fully parameterizedMtlArrayconstructors keep JACC's selected storage mode authoritative over Metal.jl's independent default.Verification
Pkg.test(): 223 passed on Apple Silicon with Julia 1.12.6 and Metal.jl 1.10.0array_storage_preference: 16 passedgit diff --check: passedReference: #402