`iterate` on `SubArray` uses slow default fallback

[jakobnissen]

It turns out the default fallback is actually rather slow for SubArray:

julia> using BenchmarkTools

julia> f(v) = sum((i for i in v), init=0);

julia> v = rand(Int, 2^12);

julia> @btime f(v);
  248.161 ns (1 allocation: 16 bytes)

julia> @btime f($(view(v, 1:2^12)));
  2.109 μs (0 allocations: 0 bytes)

julia> g(v) = sum((@inbounds v[i] for i in eachindex(v)), init=0);

julia> @btime g($(view(v, 1:2^12)));
  228.657 ns (0 allocations: 0 bytes)

I think this is rather bad as SubArray is a fairly fundamental type, and people tend to use it specifically when they are looking for performance.

It appears to be the boundschecking that is particularly expensive for SubArray. #42030 is related to this, but even if it is decided that default implementations should not use @inbounds, this should still be solved by creating a specialized method for SubArray.

This occurs on Julia 1.6.3, 1.7.0 and master as of 2021-12-02.

[N5N3]

Use a SubIterator (like ReshapedIterator for ReshapedArray) might help?

function Base.iterate(x::SubArray, state = (Iterators.product(x.indices...),))
    Base.@_propagate_inbounds_meta
    y = iterate(state...)
    y === nothing && return nothing
    x.parent[y[1]...], (state[1], Base.tail(y)...)
end

Before:

@btime f($(view(v, 1:2^12))); # 2.189 μs (0 allocations: 0 bytes)

After:

@btime f($(view(v, 1:2^12))); #1.040 μs (0 allocations: 0 bytes)

Edit: There's no performance difference if we change eltype to Float64.

[stev47]

Does the following work for you? Seems to fix it for me.

function iterate(A::SubArray, state=(eachindex(A),))
    y = iterate(state...)
    y === nothing && return nothing
    # indices of subarray are assumed valid
    @inbounds(A[y[1]]), (state[1], tail(y)...)
end

edit: On second thought, even though SubArray verifies index bounds on construction, there needs to be a guarantee that the iteration state has not been tampered with, as in #40397.
Abusing the @inbounds declaration for that contract (propagating it through iterate of SubArray as suggested above or even for AbstractArray) and adding it to all valid iteration calls e.g. in foldl might be an option.
A different approach would be for subarray to save eachindex() on construction. Then iterate would not need to trust its state and can use the @inbounds declaration.

[N5N3]

The 1-based FastContiguousSubArray should be fixed by #48720.
As for SubArray with StepRange as indices, the boundscheck is expensive as the axes is uncached for now.
Possible solution includes use static step (via Static.jl) or waiting for #47844.

I'm not sure should we close this issue. The MWE has been fixed, but the whole problem has not been solved.

[jakobnissen]

Right, let's keep the issue open. The fix in #48720 (thanks again, @matthias314) will probably fix the issue for most users, but far from all.

[matthias314]

I've started a discussion in #48773 to gauge the support for an approach based on @propagate_inbounds for iterate. Let's see what it gives.