How to simplify Array type w/o mem alloc

after reinterpreting and dropdimsing on an Array, the result is a complicated type which nominally is equivalent to just an Array. i can use collect to simplify it, but it allocates memory. is there a way to simplify it without memory allocation? here’s an MWE:

julia> using FixedPointNumbers, Colors

julia> A = rand(N0f8, 4, 300,500);

julia> @time B = dropdims(reinterpret(RGBA{N0f8}, A), dims=1);
  0.036302 seconds (219.65 k allocations: 11.421 MiB, 99.83% compilation time)

julia> @time B = dropdims(reinterpret(RGBA{N0f8}, A), dims=1);
  0.000008 seconds (2 allocations: 80 bytes)

julia> typeof(B)
Base.ReshapedArray{RGBA{N0f8}, 2, Base.ReinterpretArray{RGBA{N0f8}, 3, UInt8, Array{UInt8, 3}, false}, Tuple{}}

julia> @time C = collect(dropdims(reinterpret(RGBA{N0f8}, A), dims=1));    # only the collect is different here
  0.033743 seconds (114.19 k allocations: 6.190 MiB, 99.52% compilation time)

julia> @time C = collect(dropdims(reinterpret(RGBA{N0f8}, A), dims=1));
  0.000089 seconds (5 allocations: 608.156 KiB)

julia> typeof(C)
Matrix{RGBA{N0f8}} (alias for Array{RGBA{Normed{UInt8, 8}}, 2})

julia> B == C
true

There’s always unsafe_wrap, but as the name indicates, you need to be careful. I think it should be OK if you set it up as follows and make sure no other task has a reference to A, but don’t take my word for it.

julia> GC.@preserve A begin
           @assert size(A, 1) == 4  # sizeof(RGBA{T}) == 4sizeof(T)
           Bptr = convert(Ptr{RGBA{eltype(A)}}, pointer(A))
           B = unsafe_wrap(Array, Bptr, size(A)[2:end])
           # do stuff with B
           # but don't even think about touching A
           # for the remainder of this block
           display(B[1:10, 1:10])
           # finally, discard all unsafe references
           Bptr = nothing
           B = nothing
       end
10×10 Matrix{RGBA{N0f8}}:
 RGBA(0.714, 0.051, 0.325, 0.427)  …  RGBA(0.055, 0.824, 0.447, 0.961)
 RGBA(0.827, 0.98, 0.8, 0.553)        RGBA(0.745, 0.204, 0.125, 0.365)
 RGBA(0.4, 0.49, 0.302, 0.514)        RGBA(0.604, 0.824, 0.502, 0.161)
 RGBA(0.89, 0.616, 0.373, 0.604)      RGBA(0.573, 0.737, 0.439, 0.314)
 RGBA(0.643, 0.161, 0.165, 0.537)     RGBA(0.875, 0.133, 0.2, 0.855)
 RGBA(0.847, 0.922, 0.314, 0.565)  …  RGBA(0.835, 0.267, 0.675, 0.651)
 RGBA(0.404, 0.039, 0.471, 0.098)     RGBA(0.008, 0.714, 0.047, 0.098)
 RGBA(0.922, 0.082, 0.867, 0.576)     RGBA(0.918, 0.2, 0.31, 0.047)
 RGBA(0.808, 0.467, 0.529, 0.773)     RGBA(0.486, 0.263, 0.757, 0.639)
 RGBA(0.871, 0.216, 0.259, 0.506)     RGBA(0.031, 0.745, 0.278, 0.227)

There’s no safe way to have multiple references with different eltype to the same memory, because the compiler assumes that differently typed references never overlap (type-based aliasing analysis). Hence why this can only be achieved through unsafe_ means, putting the onus on you to stay on the straight and narrow and avoid things that would confuse the compiler or GC.

1 Like

Why not just use view?

julia> using FixedPointNumbers, Colors

julia> A = rand(N0f8, 4, 300,500);

julia> B = @btime view(reinterpret(RGBA{N0f8}, $A), 1, :, :);
  3.386 ns (0 allocations: 0 bytes)

julia> size(B)
(300, 500)

I don’t think view addresses the misgivings about having a complicated wrapper type. Allocation-wise it makes no difference; both dropdims and view are allocation-free when benchmarked properly. That said, view is a hair faster (2 ns vs. 5.5 ns on my laptop).

It doesn’t solve the wrapping issue, but you can use reinterpret(reshape, ...) to handle the dimension dropping:

julia> reinterpret(reshape, RGBA{N0f8}, A) |> typeof
Base.ReinterpretArray{RGBA{N0f8}, 2, N0f8, Base.ReinterpretArray{N0f8, 3, UInt8, Array{UInt8, 3}, false}, true}

julia> reinterpret(reshape, RGBA{N0f8}, A)
300×500 reinterpret(reshape, RGBA{N0f8}, reinterpret(N0f8, ::Array{UInt8, 3})) with eltype RGBA{N0f8}:
# ...
2 Likes