[vchuravy]

Requested by Valentin

Before: wrapping a Vector without copying (wrap(FixedSizeArray, v, dims)) gives FixedSizeArray{T,N,Vector{T}}, whose every access loads the data pointer and the Memory from the mutable Vector. In loops with conditional accesses LLVM cannot hoist these loads; under Enzyme.jl this made the augmented forward pass cache the (shadow) Memory on every inner iteration (e.g. 9.3 MB vs 1.9 MB of GC-allocated tape for a Trixi.jl flux-differencing VJP in 3D). Wrapping v.ref.mem instead is not an option, since the memory can be longer than the vector (after resize!) and the vector can start at an offset.

After: on Julia v1.11 and later, the memory backend can be a GenericMemoryRef (e.g. FixedSizeArray{T,N,MemoryRef{T}}). The array stores the reference directly and its own size gives the length, so the data pointer is a field of the immutable FixedSizeArray.

v = rand(10)
a = FixedSizeArrays.wrap(FixedSizeArray{Float64,2,MemoryRef{Float64}}, v, (2, 5))  # no copy
b = FixedSizeArrays.wrap(FixedSizeArray, Base.memoryref(v.ref.mem, 3), (2, 2))     # at an offset

How:

• The type bound of Mem is Union{DenseVector{T}, GenericMemoryRef{<:Any,T}} on v1.11+ (unchanged on v1.10). Storage access goes through small internal hooks; the reference implementations are in src/memoryref.jl and use memoryrefnew(ref, i, false) + memoryrefget/memoryrefset!/memoryref_isassigned with the default access order (so AtomicMemory works). Bounds are checked against the FixedSizeArray's size.
• wrap(FixedSizeArray, ref, size) requires at least prod(size) elements from ref on; wrap(FixedSizeArray{T,N,MemoryRef{T}} where {T,N}, v::Vector, size) wraps v.ref (exact length).
• parent returns the reference and parentindices are 1:length relative to it; dataids are those of the memory, so the array and the wrapped Vector report mightalias.
• similar, broadcasting, copy and collect_as allocate a new Memory and keep the reference backend; copyto! between references uses unsafe_copyto!; the Random and MPI extensions use a view of the memory.
• Caveat: generic code that recurses into parent(A) expecting an AbstractArray needs a method for this backend (e.g. KernelAbstractions.get_backend). As for a wrapped Vector, the array keeps the old memory alive after resize! of the vector, so it stays memory/GC safe but may no longer alias it; wrap again after resizing (documented in wrap).
• New testset "GenericMemoryRef backend" (858 tests: offsets after popfirst!, references shorter than their memory, empty arrays, String elements, aliasing with the Vector, the resize contract, GC with write barriers, overlapping copyto!, broadcasting into an aliased destination, MPI pointers, AtomicMemory, inference and no allocations), plus MemoryRef in the per-backend test loops. The full test suite passes locally on Julia 1.11.9 and 1.12.7.

In Trixi.jl (local prototype), backing the solution arrays with this backend reduces the GC-allocated Enzyme tape of an rhs! VJP from 0.67 MB to 11 KB (2D) and from 1.86 MB to 26 KB (3D) compared to Array, with identical gradients and bit-identical AMR results. Differentiating it with Enzyme.jl needs EnzymeAD/Enzyme.jl#3835 (make_zero of a MemoryRef).

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 98.49624% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.57%. Comparing base (fc90003) to head (19b939f).

Files with missing lines	Patch %	Lines
ext/FixedSizeArraysRandomExt.jl	66.66%	1 Missing ⚠️
src/memoryref.jl	98.76%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #212      +/-   ##
==========================================
+ Coverage   94.42%   95.57%   +1.14%     
==========================================
  Files           8        9       +1     
  Lines         323      429     +106     
==========================================
+ Hits          305      410     +105     
- Misses         18       19       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[giordano]

I still have to review this in details (unlikely to have much time before this evening), but in general very happy about moving in this direction!

[giordano]

Does this change the default backend?

[vchuravy]

Currently it does not. Making it the default would be a breaking change.

[giordano]

Reminder: move this entry to a different version section, when we're ready to move on with this PR

[giordano]

Caveat: generic code that recurses into parent(A) expecting an AbstractArray needs a method for this backend (e.g. KernelAbstractions.get_backend). As for a wrapped Vector, the array keeps the old memory alive after resize! of the vector, so it stays memory/GC safe but may no longer alias it; wrap again after resizing (documented in wrap).

This is probably the more problematic point of this PR, there's quite a bit of workarounds to make this kinda work.

We should also have FileCheck tests, including ensuring that memory loads are hoisted in loops with conditional access

[giordano]

This can write past the end of dst inside @inbounds:

julia> using FixedSizeArrays

julia> mem = Memory{Float64}(zeros(6));

julia> a = FixedSizeArrays.wrap(FixedSizeArray, memoryref(mem), (2,));  # the first 2 elements of `mem`

julia> b = FixedSizeArray{Float64,1,MemoryRef{Float64}}(ones(4));

julia> f!(a, b) = @inbounds copyto!(a, b)
f! (generic function with 1 method)

julia> f!(a, b)
2-element FixedSizeVector{Float64, MemoryRef{Float64}}:
 1.0
 1.0

julia> mem
6-element Memory{Float64}:
 1.0
 1.0
 1.0
 1.0
 0.0
 0.0

[giordano]

Should there be a method to wrap a Memory (with memoryrefnew(mem))?

[giordano]

Should this be aware of the offsets, so that non-overlapping slices of the same memory don't alias?

[giordano]

This means that the backend type may not be preserved?

[giordano]

Do we still need this method at all with MPI.jl@0.20.28?

[giordano]

The check contradicts the comment that this should apply to GenericMemoryRef?