JuliaArrays/LazyArrays.jl

Indexing performance for heterogeneous concatenated arrays

jipolanco opened this issue · 1 comments

Indexing heterogeneous concatenated arrays is slow due to type instability issues. By heterogeneous I mean something like:

A = Vcat(
    rand(1:10, 12),
    @SVector(zeros(Int, 4)),
    4:42,
)

where each concatenated subarray may have different type.

Indexing the first subarray is relatively fast, while indexing subsequent arrays becomes slower and allocates memory. On Julia 1.5.1:

N = length(A)
@btime $A[1];   # 7.356 ns (0 allocations: 0 bytes)
@btime $A[14];  # 43.875 ns (2 allocations: 112 bytes)
@btime $A[$N];  # 80.788 ns (4 allocations: 208 bytes)

This seems to be caused by type instabilities in vcat_getindex, as can be checked with @code_warntype LazyArrays.vcat_getindex(A, N). More precisely, the problem is in the loop:

    for A in f.args
        n = length(A)
        κ  n && return convert(T,A[κ])::T
        κ -= n
    end

where A can be any of the concatenated subarrays.

The same issue occurs with setindex!, and also with two-dimensional concatenated arrays (Vcat{2} and Hcat).

I've been playing with a possible fix using recursive iteration over the f.args tuple, and I'll submit a PR soon.

Closed by #133