ArrayInterface.jl icon indicating copy to clipboard operation
ArrayInterface.jl copied to clipboard

Layouts - v2

Open Tokazama opened this issue 4 years ago • 2 comments
trafficstars

Following feedback from #141 I've rewritten the approach to layouts. The basic interface is as follows:

Given the array type ArrayType the layout_plan method may be written as follows.

function layout_plan(::Type{A}, s::AccessStyle) where {A<:ArrayType}
    # if `ArrayType` is the final index conversion
    ArrayTypeIndex

    # if `ArrayType` transforms whatever index it is given and passes it on to subsequent layers
    # `combine_layouts` takes whatever nested layers produces and provides an opportunity to combine it
    combine_layouts(A, layout_plan(parent_type(A), s))

    # if `s` is not a compatible access style propagate a new one and provide a converting index
    LayoutPlan(ConvertArrayTypeIndex, layout_plan(A, ProperAccess()))
end
Short note on implementation

LayoutPlan was inspired by Base.Broadcast.Broadcasted and a bit by AbstractFFts.Plan. Where Broadcasted flattens out and combines features and AbstractFFts.Plan uses type information to construct a new method based on an array type. I've used subtypes of ArrayIndex instead of traits so that we can use them as method calls and so they can hold dynamic information. This is a bit awkward when dealing with more complicated types such as StrideIndex but it's also a lot simpler than creating a method, trait, and index type for each type of index.

It's arguably easier to construct layouts with actual instances (e.g., layout(x::X, s::AccessStyle) = combine_layouts(x, layout(parent(x), s))), but I tried that and it completely kills performance. However, constructing a layout plan first with types also has the benefit of freeing us from worrying about constant propagation and potentially adding in other passes over the LayoutPlan before materializing it.

Benchmarks

There's a lot of clean up that still needs to happen but initial benchmarks appear adequate.

a = Array{Int}(undef, 5, 2, 5, 2); 
a[:] .= 1:100;
av = view(a, :, 1, :, :);
aperm = PermutedDimsArray(a, (3,1,4,2));
mv = view(aperm, :, 1, :, 1);
mvp = mv';

julia> @btime ArrayInterface.getindex(a, :, 1, :, 1);
  492.411 ns (3 allocations: 368 bytes)

julia> @btime getindex(a, :, 1, :, 1);
  504.557 ns (10 allocations: 592 bytes)

julia> @btime ArrayInterface.getindex(aperm, :, 1, :, 1);
  490.788 ns (3 allocations: 240 bytes)

julia> @btime getindex(aperm, :, 1, :, 1);
  494.326 ns (13 allocations: 512 bytes)

julia> @btime ArrayInterface.getindex(av, :, 1, :);
  449.853 ns (3 allocations: 208 bytes)

julia> @btime getindex(av, :, 1, :);
  412.600 ns (9 allocations: 432 bytes)

julia> @btime ArrayInterface.getindex(mv, :, 1:2);
  221.106 ns (6 allocations: 384 bytes)

julia> @btime getindex(mv, :, 1:2);
  232.407 ns (3 allocations: 224 bytes)

julia> @btime ArrayInterface.getindex(mvp, :, 1:2);
  200.144 ns (6 allocations: 336 bytes)

julia> @btime getindex(mvp, :, 1:2);
  243.843 ns (3 allocations: 176 bytes)

Indexing SubArrays was consistently slower (as shown in these benchmarks), but I've yet to optimize indexing with layouts. I'm actually surprised that performance is consistently doing better given the same exact array types and lack of any explicit specialization in constructing loops.

For Review

  • Biggest concern is whether the general interface makes sense. I'll certainly write documentation with more explicit examples than in this comment, but it's hard to know if if this makes sense in general. Months of coming back to this may have distorted my view.
  • refdata - It might make more sense to have ManualMemory own this behavior. I don't care what the name is, maybe preserve_buffer
  • reinterpret layout - this would require a unique buffer for reinterpreted values and probably a special index type. something else that could be in ManualMemory.jl?

Once we have layouts in place I think the only unique component of the array interface that still needs to be available here is constructing new instances of an array. Assuming ManualMemory.jl becomes responsible for creating buffers, we could probably have a fairly simple method for wrapping buffers around newly constructed axes and attaching relevant metadata.

Tokazama avatar Jun 07 '21 20:06 Tokazama

I'm going to do a slight refractor of layout_plan to resolve world age issues.

Tokazama avatar Jun 08 '21 07:06 Tokazama

I'm going to do a slight refractor of layout_plan to resolve world age issues.

I'll have to look at this later, but on the subject of world age issues, could you also make sure tests pass with --compiled-modules=no? MPI.jl suggests that as one of two options to work around a parallel precompilation issue. I think people really should be using the first of those two options instead, but the fact of the matter is I've gotten a couple issues on other packages from people taking the second approach, so it'd be good for that to be supported.

chriselrod avatar Jun 08 '21 08:06 chriselrod

Anything static should go to https://github.com/JuliaArrays/StaticArrayInterface.jl

ChrisRackauckas avatar Feb 18 '23 11:02 ChrisRackauckas