JuliaLang / julia

The Julia Programming Language
https://julialang.org/
MIT License
45.42k stars 5.45k forks source link

`stack` is slow for certain iterators #52590

Open thchr opened 8 months ago

thchr commented 8 months ago

I noticed that stack is slower than a naive reduce-based implementation for nested SVectors:

using StaticArrays, BenchmarkTools

rstack(iter) = reduce(hcat, iter)
iter = SVector{3, SVector{3, Float64}}(SVector{3, Float64}(1.0, 0.0, 0.0), 
                                       SVector{3, Float64}(-0.5, 0.8660254037844387, 0.0), 
                                       SVector{3, Float64}(0.0, 0.0, 0.7923832584290136))

@btime stack($iter);  # 26.788 ns (1 allocation: 80 bytes)
@btime rstack($iter); #  1.800 ns (0 allocations: 0 bytes)

rstack(iter) == stack(iter) # true

The code path seems to be via _typed_stack(::Colon, ::Float64, ::SVector{3, Float64}, iter). I could not immediately figure out what is causing the slowdown.

mcabbott commented 8 months ago

stack(iter) isa MMatrix, made I believe by similar and then the same code paths used for Arrays.

But @which reduce(hcat, iter) points to a method reduce(::typeof(hcat), A::StaticArray{<:Tuple,<:StaticVecOrMatLike}) in StaticArrays, which makes an SMatrix. That package could define similar fast methods for stack acting on its types.