Improve performance of `eachslice`

[oxinabox]

We should add eachslice
in julia 1.1 this can be an overload for Base.eachslice
and in earlier we can just export it.

Right now can kinda fake it via; vec(mapslices(nda; dim=dim))

Example

outputs = NamedDimsArray{(:variates, :observations)}(outputs)
table = mapslices(outputs, dims=:variates) do row
        NamedTuple{colnames, coltypes}(Tuple(row))
end |> vec

[nickrobinson251]

In 1.1, is this sufficient

julia> Base.eachslice(A::NamedDimsArray{L}; dims) where L = eachslice(parent(A); dims=NamedDims.dim(L, dims))

julia> nda = NamedDimsArray(rand(2, 2, 2), (:a, :b, :c));

julia> g = eachslice(nda; dims=:c);

julia> iterate(g)
([0.20502 0.326133; 0.33411 0.0193653], 1)

julia> typeof.(iterate(g))
(SubArray{Float64,2,Array{Float64,3},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},Int64},true}, Int64)

Or do we want to have the SubArrays be NamedDimsArrays too?

[oxinabox]

We want the slice to have names too.
As a rule if a thing can have anamw it should have a name.

But that is not hard, I don't think.
Use Base.Generator to do the wrapping.

Trick to me is working out how to do it in 1.0.
Which I think requires fixing Compat.jl

[nickrobinson251]

Fixing Compat.jl is the better way to do this...

But if that's loads of work, we could just copy the Base.eachslice definition here?

[nickrobinson251]

I do not know how to get allocations down on something like this

function Base.eachslice(a::NamedDimsArray; dims, kwargs...)
     numerical_dims = dim(a, dims)
     data = Base.eachslice(parent(a); dims=numerical_dims, kwargs...)
     new_names = slice_names(a, numerical_dims)
     return Base.Generator(NamedDimsArray{new_names}, data)
end

function slice_names(a::NamedDimsArray{L}, n::Integer) where L
     return compile_time_return_hack((L[1:n-1]..., L[n+1:end]...))
end

julia> a = rand(30, 30, 30);

julia> nda = NamedDimsArray(a, (:a, :b, :c));

julia> @btime (() -> eachslice($a; dims=3))();
  185.928 ns (3 allocations: 64 bytes)

julia> @btime (() -> eachslice($nda; dims=:c))();
  4.917 μs (13 allocations: 544 bytes)

I guess because the generator returned by the inner eachslice gets collected when the final Generator is constructed

I am not sure where the allocations come from...

julia> @btime Base.Generator(identity, eachslice($a; dims=3));
  257.180 ns (4 allocations: 80 bytes)

But slice_names needs improving

julia> @btime (() -> NamedDims.slice_names($nda, 1))();
  3.659 μs (3 allocations: 160 bytes)

[oxinabox]

You should be able to use remaining_dimnames_after_dropping
Which we already have defined

Rather than your new slice_names

[nickrobinson251]

wow - very useful - and i was not about to figure out the hack that makes that function not allocate

thanks

[nickrobinson251]

but unfortunately swapping in remaining_dimnames_after_dropping doesn't change the amount eachslice allocated... :(

function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
    numerical_dims = dim(a, dims)
    data = eachslice(parent(a); dims=numerical_dims, kwargs...)
    new_names = remaining_dimnames_after_dropping(L, numerical_dims)
    return Base.Generator(NamedDimsArray{new_names}, data)
end

julia> @btime (() -> eachslice($nda; dims=:c))();
  4.258 μs (13 allocations: 560 bytes)

[nickrobinson251]

Here's where the allocations come from

# `eachslice` by itself allocates 3 times
julia> function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
           numerical_dims = NamedDims.dim(a, dims)
           data = Base.eachslice(parent(a); dims=numerical_dims, kwargs...)

           return data
       end

julia> @btime (() -> eachslice($nda; dims=:c))();
  199.451 ns (3 allocations: 64 bytes)

# somehow computing `new_names` adds 3
julia> function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
           numerical_dims = NamedDims.dim(a, dims)
           data = Base.eachslice(parent(a); dims=numerical_dims, kwargs...)
           new_names = NamedDims.remaining_dimnames_after_dropping(L, numerical_dims)
           return data
       end

julia> @btime (() -> eachslice($nda; dims=:c))();
  3.111 μs (6 allocations: 240 bytes)

# the `Generator` call adds 1
julia> function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
           numerical_dims = NamedDims.dim(a, dims)
           data = Base.eachslice(parent(a); dims=numerical_dims, kwargs...)
           new_names = NamedDims.remaining_dimnames_after_dropping(L, numerical_dims)
           return Base.Generator(identity, data)
       end

julia> @btime (() -> eachslice($nda; dims=:c))();
  3.191 μs (7 allocations: 256 bytes)

# passing the `NamedDimsArray` constructor adds 6
julia> function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
           numerical_dims = NamedDims.dim(a, dims)
           data = Base.eachslice(parent(a); dims=numerical_dims, kwargs...)
           new_names = NamedDims.remaining_dimnames_after_dropping(L, numerical_dims)
           return Base.Generator(NamedDimsArray{new_names}, data)
       end

julia> @btime (() -> eachslice($nda; dims=:c))();
  4.143 μs (13 allocations: 560 bytes)

[oxinabox]

Got an order of magnitude improvement,
but still twice as slow and as allocating as it should be

function Base.eachslice(a::NamedDimsArray{L}; dims, kwargs...) where L
    numerical_dims = dim(a, dims)
    slices = eachslice(parent(a); dims=numerical_dims, kwargs...)
    
    return Base.Generator(slices) do sl
        # For unknown reasons (something to do with hoisting?) having this in the function
        # actually results in less memory being allocated 
        new_names = remaining_dimnames_after_dropping(L, numerical_dims)
        return NamedDimsArray(sl, new_names)
    end 
end

nda = NamedDimsArray(rand(2, 2, 2), (:a, :b, :c,));

@btime (()->eachslice($nda; dims=:c))()
@btime eachslice($nda; dims=:c)
# 258.457 ns (5 allocations: 112 bytes)

[mcabbott]

I had a look at this. Why define a custom iterator for eachslice at all, but not for eachcol? If I simply comment out that definition, then it’s much faster, here’s a test:

ab = NamedDimsArray((1:3) .+ zeros(Int,3)', (:a, :b));
@btime [sum(x)^2 for x in eachcol($ab)] #  78.411 ns (7 allocations: 320 bytes)

# on master (same with dims=2)
@btime [sum(x)^2 for x in eachslice($ab, dims=:b)] # 13.842 μs (40 allocations: 1.73 KiB)

# with eachslice commented out:
@btime [sum(x)^2 for x in eachslice($ab, dims=2)] # 378.542 ns (8 allocations: 336 bytes)

In both cases, [NamedDims.names(x) for x in eachslice(ab, dims=2)] and [NamedDims.names(x) for x in eachcol(ab)] see the name :a, thanks to view. Thus I think the only thing needed is to translate dims=:b to a numerical dimension. Do you want a PR?

[oxinabox]

@nickrobinson251 is tracking this better than I am

[nickrobinson251]

I'm not sure I follow completely... but absolutely make a PR if you've a way to get the same eachslice(ab, dims=:b) behaviour only faster!