Add an opt-in Metal array storage preference

docs/src/api_usage.md

@@ -10,7 +10,7 @@ JACC APIs consist of three main components:
 
 ## Backend selection
 
-- **`JACC.set_backend`**: allows selecting the runtime backend on **CPU**: `Threads` (default) and **GPU**: `CUDA`, `AMDGPU`, `oneAPI`. Uses Preferences.jl and stores the selected backend in a [LocalPreferences.jl](https://github.com/JuliaPackaging/Preferences.jl) file if JACC.jl is a project dependency. Use `JACC.set_backend` prior to running any code targeting a particular backend.
+- **`JACC.set_backend`**: allows selecting the runtime backend on **CPU**: `Threads` (default) and **GPU**: `CUDA`, `AMDGPU`, `oneAPI`, `Metal`. Uses Preferences.jl and stores the selected backend in a [LocalPreferences.toml](https://github.com/JuliaPackaging/Preferences.jl) file if JACC.jl is a project dependency. Use `JACC.set_backend` prior to running any code targeting a particular backend.
 
 Example:
 ```julia
@@ -54,7 +54,7 @@ JACC.@init_backend
 
 ## Memory allocation
 
-- **`JACC.array()`**: create a new array on the device with the specified type and size.
+- **`JACC.array()`**: create a new array on the device with the specified type and size, or copy a host array to the device. Metal projects can opt into unified (`SharedStorage`) memory for host-array copies with `JACC.set_backend("Metal"; storage = :shared)`; portable code continues to call `JACC.array(x)` without backend-specific keywords. Host-array copies default to `:private`. Shared storage avoids the private-storage transfer path; its construction-latency and GPU-bandwidth tradeoffs should be measured for the target workload.
 - **`JACC.zeros`**: create a new array on the device filled with zeros.
 - **`JACC.ones`**: create a new array on the device filled with ones.
 - **`JACC.fill`**: create a new array on the device filled with a specified value.
@@ -202,4 +202,4 @@ where the additional parameters are:
 - `stream`: stream identifier (GPU only), handler from `JACC.default_stream()` or `JACC.create_stream()` see [AMD GPU tests](https://github.com/JuliaGPU/JACC.jl/blob/main/test/backend/amdgpu.jl)
 - `sync`: true or false, whether to synchronize after kernel launch (default: true)
 
-Other examples and more advanced usages can be found in the [JACC tests directory](https://github.com/JuliaGPU/JACC.jl/blob/main/test/unittests.jl)
+Other examples and more advanced usages can be found in the [JACC tests directory](https://github.com/JuliaGPU/JACC.jl/blob/main/test/unittests.jl)

ext/MetalExt/MetalExt.jl

@@ -409,6 +409,49 @@ JACC.sync_workgroup(::MetalBackend) = Metal.threadgroup_barrier()
 
 JACC.array_type(::MetalBackend) = Metal.MtlArray
 
-JACC.array(::MetalBackend, x::Base.Array) = Metal.MtlArray(x)
+function _compute_array_storage()
+    preferences = get(
+        JACC.Preferences.Backend._EXT_PREFS[], "metal", Dict{Symbol, Any}())
+    value = get(preferences, :storage, "private")
+    value isa Union{AbstractString, Symbol} || throw(ArgumentError(
+        "Invalid Metal array storage: $(repr(value)); " *
+        "expected :private or :shared"))
+    storage = lowercase(String(value))
+    storage in ("private", "shared") || throw(ArgumentError(
+        "Invalid Metal array storage: $(repr(value)); " *
+        "expected :private or :shared"))
+    return storage
+end
+
+# Cached module global rather than a plain module-load-once value: `storage`
+# is one of the extension preferences `set_backend` is documented (and
+# tested, see array_storage_preference in test/backend/metal.jl) to apply
+# live within the same Julia session, without a restart. _EXT_PREFS_GENERATION
+# is bumped on every write, so this stays a cheap Int comparison on the
+# array-allocation hot path instead of a Dict lookup plus revalidation on
+# every call, while still tracking live changes.
+const _ARRAY_STORAGE_CACHE = Ref{Tuple{Int, String}}((-1, ""))
+
+function _array_storage()
+    generation = JACC.Preferences.Backend._EXT_PREFS_GENERATION[]
+    cached_generation, cached_value = _ARRAY_STORAGE_CACHE[]
+    if cached_generation == generation
+        return cached_value
+    end
+    value = _compute_array_storage()
+    _ARRAY_STORAGE_CACHE[] = (generation, value)
+    return value
+end
+
+function JACC._array(::MetalBackend, x::AbstractArray)
+    if _array_storage() == "shared"
+        return Metal.MtlArray{eltype(x), ndims(x), Metal.SharedStorage}(x)
+    end
+    return Metal.MtlArray{eltype(x), ndims(x), Metal.PrivateStorage}(x)
+end
+
+JACC._array(::MetalBackend, x::Metal.MtlArray) = x
+JACC.array(backend::MetalBackend, x::Base.Array) = JACC._array(backend, x)
+JACC.array(::MetalBackend, x::Metal.MtlArray) = x
 
 end # module MetalExt

src/array.jl

@@ -71,10 +71,18 @@ end
 
 """
     array([T=default_float()], dims...)
+    array(x::AbstractArray)
 
-Create an uninitialized array on the device with the specified type and size.
+Create an uninitialized array on the device with the specified type and size,
+or copy the host array `x` to the device.
+
+Backend-specific allocation policy is configured outside portable code. For
+example, `set_backend("Metal"; storage = :shared)` makes host-array copies use
+Metal unified memory for that project. Host-array copies default to device-private.
 """
-array(x::AbstractArray) = to_device(x)
+array(x::AbstractArray) = _array(default_backend(), x)
+_array(::Any, x::AbstractArray) = to_device(x)
+
 array(::Type{T}, dims) where {T} = array_type(){T, length(dims)}(undef, dims)
 array(::Type{T}, dims...) where {T} = array(T, dims)
 array(dims) = array(default_float(), dims)

src/preferences.jl

@@ -91,6 +91,33 @@ const _DEFAULT = Ref(String(default))
 const list = @load_preference("backends", ["threads"])
 const _LIST = Ref(deepcopy(list))
 const _PLACE = Ref(@load_preference("placement", Dict{String, String}()))
+const extension_preferences = @load_preference(
+    "extension_preferences", Dict{String, Any}())
+
+_serialize_extension_preference(value::Symbol) = String(value)
+function _serialize_extension_preference(value::AbstractDict)
+    return Dict(
+        String(key) => _serialize_extension_preference(item)
+        for (key, item) in value)
+end
+function _serialize_extension_preference(value::Union{Tuple, AbstractVector})
+    return [_serialize_extension_preference(item) for item in value]
+end
+_serialize_extension_preference(value) = value
+
+function _runtime_extension_preferences(preferences)
+    return Dict{String, Dict{Symbol, Any}}(
+        String(backend) => Dict{Symbol, Any}(
+            Symbol(key) => value
+            for (key, value) in values)
+        for (backend, values) in preferences)
+end
+
+const _EXT_PREFS = Ref(_runtime_extension_preferences(extension_preferences))
+# Bumped on every mutation of _EXT_PREFS so extensions can cheaply detect
+# staleness of any value they cache from it, without recomputing on every
+# call. See ext/MetalExt/MetalExt.jl `_array_storage` for the reader side.
+const _EXT_PREFS_GENERATION = Ref(0)
 
 const package_names = ["CUDA", "AMDGPU", "oneAPI", "Metal"]
 
@@ -140,6 +167,9 @@ function unset_backend()
     @delete_preferences!("default_backend")
     @delete_preferences!("backends")
     @delete_preferences!("placement")
+    @delete_preferences!("extension_preferences")
+    empty!(Preferences.Backend._EXT_PREFS[])
+    Preferences.Backend._EXT_PREFS_GENERATION[] += 1
     @info """
         Backend preferences deleted
         Restart your Julia session for this change to take effect!
@@ -171,20 +201,37 @@ function set_default_backend(new_backend::Symbol)
     set_default_backend(String(new_backend))
 end
 
-function set_backend(b::AbstractString)
+function _set_extension_preferences(backend::String, kw)
+    values = Dict{Symbol, Any}(pairs((; kw...)))
+    isempty(values) && return
+
+    preferences = deepcopy(Preferences.Backend._EXT_PREFS[])
+    preferences[backend] = values
+    serialize = Preferences.Backend._serialize_extension_preference
+    persisted = Dict(
+        name => serialize(settings)
+        for (name, settings) in preferences)
+    @set_preferences!("extension_preferences"=>persisted)
+    Preferences.Backend._EXT_PREFS[] = preferences
+    Preferences.Backend._EXT_PREFS_GENERATION[] += 1
+end
+
+function set_backend(b::AbstractString; kw...)
     nb = lowercase(b)
     if Preferences.Backend._LIST[] == [nb]
         if Preferences.Backend._DEFAULT[] == nb
+            _set_extension_preferences(nb, kw)
             return
         end
     else
         _check_supported(nb)
         unset_backend()
     end
     set_default_backend(nb)
+    _set_extension_preferences(nb, kw)
 end
 
-set_backend(b::Symbol) = set_backend(String(b))
+set_backend(b::Symbol; kw...) = set_backend(String(b); kw...)
 
 function add_backend(new_backend::AbstractString)
     new_backend_lc = lowercase(new_backend)

test/backend/metal.jl

@@ -21,6 +21,45 @@ end
     @test eltype(x) == Float32
 end
 
+@testset "array_storage_preference" begin
+    using Metal
+    using Suppressor
+    N = 10
+    h = ones(Float32, N)
+    x = JACC.array(h)
+    @test typeof(x) == MtlArray{Float32, 1, Metal.PrivateStorage}
+    @test JACC.array(x) === x
+    @test JACC.array(JACC.default_backend(), x) === x
+    @test JACC.array(JACC.default_backend(), h) isa
+          MtlArray{Float32, 1, Metal.PrivateStorage}
+    try
+        @suppress JACC.set_backend("Metal"; storage = :shared)
+        preferences = load_preference(JACC, "extension_preferences")
+        @test preferences["metal"]["storage"] == "shared"
+        @test JACC.Preferences.Backend._EXT_PREFS[]["metal"] ==
+              Dict(:storage => :shared)
+        xs = JACC.array(h)
+        @test typeof(xs) == MtlArray{Float32, 1, Metal.SharedStorage}
+        @test Metal.is_shared(xs)
+        @test Array(xs) == h
+        @test JACC.array(xs) === xs
+        @test JACC.array(JACC.default_backend(), xs) === xs
+        @test JACC.array(JACC.default_backend(), h) isa
+              MtlArray{Float32, 1, Metal.SharedStorage}
+        h2 = ones(Float32, N, N)
+        xs2 = JACC.array(h2)
+        @test typeof(xs2) == MtlArray{Float32, 2, Metal.SharedStorage}
+    finally
+        @suppress JACC.set_backend("Metal"; storage = :private)
+    end
+    @test JACC.array(h) isa MtlArray{Float32, 1, Metal.PrivateStorage}
+    @suppress JACC.set_backend("Metal"; storage = :bogus)
+    @test_throws ArgumentError JACC.array(h)
+    @suppress JACC.set_backend("Metal"; storage = 1)
+    @test_throws ArgumentError JACC.array(h)
+    @suppress JACC.set_backend("Metal"; storage = :private)
+end
+
 include("preferences.jl")
 
 @testset "preferences" begin

test/backend/preferences.jl

@@ -18,12 +18,20 @@ function test_preferences(bksym::Symbol)
         @test isempty(JACC.Preferences.Backend._PLACE[])
     end
 
+    @suppress JACC.set_backend(bkstr; test_preference = true)
+    extension_preferences = load_preference(JACC, "extension_preferences")
+    @test extension_preferences[bkstr]["test_preference"]
+    @test JACC.Preferences.Backend._EXT_PREFS[][bkstr] ==
+          Dict(:test_preference => true)
+
     # Clear settings
     @suppress JACC.unset_backend()
     @test isempty(JACC.Preferences.Backend._LIST[])
     @test isempty(JACC.Preferences.Backend._PLACE[])
     @test load_preference(JACC, "backends") == nothing
     @test load_preference(JACC, "default_backend") == nothing
+    @test load_preference(JACC, "extension_preferences") == nothing
+    @test isempty(JACC.Preferences.Backend._EXT_PREFS[])
 
     # "not a backend"
     @test_throws ArgumentError JACC.set_backend("NAB")

test/backend/threads.jl

@@ -52,6 +52,10 @@ end
     @test typeof(x) == Array{Complex{Float32}, 3}
 end
 
+@testset "array_storage_unsupported" begin
+    @test_throws MethodError JACC.array(ones(Float32, 10); storage = :shared)
+end
+
 include("preferences.jl")
 
 @testset "preferences" begin

test/unittests.jl

@@ -59,6 +59,11 @@ end
     @test ndims(x) == 3
     @test eltype(x) == Complex{Float32}
     @test size(x) == (5, 5, 5)
+
+    # Copy from host using the backend's configured storage policy.
+    h = ones(Float32, 10)
+    x = JACC.array(h)
+    @test JACC.to_host(x) == h
 end
 
 @testset "transfer!" begin