From b4ecf2776458b648f976cb4d1f607f99e096885c Mon Sep 17 00:00:00 2001 From: PraneethMerugu Date: Wed, 9 Sep 2026 15:49:27 -0400 Subject: [PATCH] Copy allocated arrays through backend-owned transfers --- docs/src/api/localmath.md | 3 ++- src/bound_law.jl | 17 +++++++++++++++-- test/metal/localmath_authoring.jl | 17 ++++++++++++++--- test/runtests.jl | 1 + test/test_storage_authoring.jl | 9 +++++++++ 5 files changed, 41 insertions(+), 6 deletions(-) diff --git a/docs/src/api/localmath.md b/docs/src/api/localmath.md index 79e4ea9..aaca8c1 100644 --- a/docs/src/api/localmath.md +++ b/docs/src/api/localmath.md @@ -39,7 +39,8 @@ output_array = LocalMath.storage(prepared, output) ``` `allocate(value)` fills a Field when `value` has its exact element type; -`allocate(source)` copies an exact-shape source array to independent backend +`allocate(source)` copies an exact-shape source array, including an ordinary +host array view on the qualified CPU and Metal paths, to independent backend storage; and `allocate()` creates the exact bounded storage for a produced Collection. A caller-owned `StructArray` is borrowed unchanged. Allocating one copies its component arrays recursively and preserves the record layout. diff --git a/src/bound_law.jl b/src/bound_law.jl index a38d273..7f6a2d6 100644 --- a/src/bound_law.jl +++ b/src/bound_law.jl @@ -77,8 +77,21 @@ end function _copy_allocated_array(backend, source::AbstractArray) destination = _allocate_array(backend, eltype(source), size(source)) - copyto!(destination, source) - KernelAbstractions.synchronize(backend) + if backend isa KernelAbstractions.CPU + copyto!(destination, source) + return destination + end + # Backend copies intentionally accept a narrower set of physical arrays + # than LocalMath's public Allocate contract. Materialize an unsupported + # host view before crossing that boundary; supported device-to-device and + # dense host transfers retain the backend's native copy path. + transfer_source = applicable( + KernelAbstractions.copyto!, backend, destination, source + ) ? source : Array(Adapt.adapt(KernelAbstractions.CPU(), source)) + GC.@preserve destination transfer_source begin + KernelAbstractions.copyto!(backend, destination, transfer_source) + KernelAbstractions.synchronize(backend) + end return destination end diff --git a/test/metal/localmath_authoring.jl b/test/metal/localmath_authoring.jl index c09d535..5761f98 100644 --- a/test/metal/localmath_authoring.jl +++ b/test/metal/localmath_authoring.jl @@ -30,16 +30,27 @@ struct LocalMathMetalNode end static_law = LocalMath.@localmath i ∈ cells begin static_output[i] = static_input[i] end - static_source = StaticVector[ - StaticVector(Float32(i), Float32(i + 1)) for i in 1:4] + static_source_parent = StaticVector[ + StaticVector(Float32(i), Float32(i + 1)) for i in 1:8] + static_source_view = @view static_source_parent[1:2:7] + static_source = collect(static_source_view) # This case qualifies cold structured allocation only; SVector field # execution is intentionally outside the reviewed Metal storage operations. static_bound = LocalMath.bind(static_law, - static_input => LocalMath.Allocate(static_source), + static_input => LocalMath.Allocate(static_source_view), static_output => LocalMath.Allocate(undef); backend) + static_source_parent[1] = StaticVector(-1.0f0, -2.0f0) @test Array(LocalMath.storage(static_bound, static_input)) == static_source + @test LocalMath.storage(static_bound, static_input) !== static_source_view @test size(LocalMath.storage(static_bound, static_output)) == (4,) + device_source = Metal.MtlArray(static_source) + device_bound = LocalMath.bind(static_law, + static_input => LocalMath.Allocate(device_source), + static_output => LocalMath.Allocate(undef); + backend) + @test Array(LocalMath.storage(device_bound, static_input)) == static_source + @test LocalMath.storage(device_bound, static_input) !== device_source source = LocalMath.Space(LocalMathMetalNode, 3) destination = LocalMath.Space(LocalMathMetalNode, 2) diff --git a/test/runtests.jl b/test/runtests.jl index 1ecc964..53b8389 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -74,6 +74,7 @@ const LOCALMATH_TEST_SUITE = Dict{String, Expr}( :apply_type, :checked_mul, :code_typed_by_type, + :copyto!, :datatype_alignment, :device, :functional, diff --git a/test/test_storage_authoring.jl b/test/test_storage_authoring.jl index f03fc50..229cd66 100644 --- a/test/test_storage_authoring.jl +++ b/test/test_storage_authoring.jl @@ -47,6 +47,15 @@ end source[1] = 99f0 @test input_storage[1] == 1f0 + source_view = @view source[2:-1:1, :] + view_bound = LMA.bind(law, + input => LMA.Allocate(source_view), + output => LMA.Allocate(undef); backend) + @test LMA.storage(view_bound, input) == source_view + @test LMA.storage(view_bound, input) !== source_view + source[2] = 88f0 + @test LMA.storage(view_bound, input)[1] == 2f0 + prepared = LMA.prepare(law, input => source, output => LMA.Allocate(undef); backend) @test LMA.storage(prepared.plan, input) === source