Tree traversal really slow due to ridiculous ammounts of compilation

[oscardssmith]

If we define the simplest possible tree

using AbstractTrees
struct Tree
    child :: Union{Tree, Int}
end
AbstractTrees.children(tree::Tree) = (tree.child, )
function nest(n)
    tree = Tree(1)
    for i in 1:n
        tree = Tree(tree)
    end
    tree
end
function it(x)
    for x in PreOrderDFS(tree)
        x isa Int && return x
    end
end

Then the current definition of PreOrderDFS (and other traversals) recompile for different sized trees.

julia> tree = nest(2);
julia> @time it(tree);
  0.110824 seconds (297.39 k allocations: 19.659 MiB, 11.69% gc time, 99.77% compilation time)

julia> tree = nest(100);
julia> @time it(tree);
 20.828098 seconds (4.94 M allocations: 329.490 MiB, 0.25% gc time, 99.96% compilation time)

Defining AbstractTrees.childtype(::Tree) = Union{Int, Tree} makes it roughly 50x faster, but there still is compile time proportional to the size of tree.

[oscardssmith]

Note that this is a major regression. In 0.3.4 @ExpandingMan, the same definitions give

julia> @time it(tree);
  0.194956 seconds (705.19 k allocations: 46.591 MiB, 5.54% gc time, 99.85% compilation time)

julia> tree = nest(100);

julia> @time it(tree);
  0.000387 seconds (3.73 k allocations: 917.219 KiB)

[ExpandingMan]

It seems that what's happening is that the cursor is recursively wrapping itself so that it has to recompile for every iteration:

◖◗ s = AbstractTrees.PreOrderState(t);

◖◗ typeof(s |> next)
AbstractTrees.PreOrderState{ImplicitCursor{ImplicitCursor{Nothing, Tree, AbstractTrees.InitialState}, Tree, Int64}}

◖◗ typeof(s |> next |> next)
AbstractTrees.PreOrderState{ImplicitCursor{ImplicitCursor{ImplicitCursor{Nothing, Tree, AbstractTrees.InitialState}, Tree, Int64}, Tree, Int64}}

It's easy to see why this evaded my sanity checks because I do not normally benchmark compile times.

This is pretty scary because as I recall not having to deal with unwrapping cursors when they are passed to each other was a significant simplification.

I'll look into it. I was probably too aggressive about trying to make things type stable. The right position to take here is probably that either you have a tree with all nodes of the same type or tough luck, all bets are off.

I don't think this will be an easy fix so any help is welcome.

[oscardssmith]

yeah. I only found it because it stopped code I was using from terminating (at least in less than a minute).

[ExpandingMan]

It looks like in this particular case it should be rather easy to fix: those recursively wrapped cursors aren't doing anything, just unwrap them it doesn't matter. What has me worried is that I feel rather unsure of the broader implications of doing that, it might make it extremely hard for tree traversal to be type stable as once you have done this cursors cannot know the exact types of their neighbors.

I have a bad feeling that the result of this is pretty much going to be ripping the guts out of the package and saying to users "Either define NodeType or get shit performance 🤷 ". This is definitely a particularly tricky use case for Julia's type system.

[ExpandingMan]

I've looked into this further.

It seems that the crux of the problem is the need to store the iteration state of the parent. Right now it is "cheating" by storing the entire tree along with the states within each cursor, but for this to be type stable requires the recursive typing of the cursors we see above. The only solution the the dilemma that I can see would be to loosen the requirement that parent and nextsibling from a cursor are type stable.

By itself, this isn't so bad, again because of the nature of trees it's not reasonable to ask for traversal over them to be type stable without providing a guarantee of the type of all connected nodes. Unfortunately, if I do this naively, I will make it impossible for tree iteration to ever be type stable. It's also going to be incredibly difficult to make sure that it is type stable in cases where nodetype has been defined but the iteration state type is not guaranteed. If we cannot even guarantee type stability in all cases where nodetype is defined it will cripple the package.

So, it seems I've gotten myself in quite a hole here. In a way it seems crazy that I did not notice this problem initially, but as I am so accustomed to discounting work required of the compiler it's also not very hard to see why I went this route: it is not only simpler than the alternatives but gives us everything else we want at a cost to only the compiler.

I am committed to fixing this as I consider this to be an important package in the ecosystem and the current situation is entirely my fault, but it might take me a week or more to come to a satisfactory conclusion. Trying to do it without breaking changes is also challenging (even changing the type signatures of the cursors is technically breaking).

[oscardssmith]

I'll work on incorporating https://gist.github.com/YingboMa/acd8dea38e708a4d6a127b58709cc487 which should fix the issue for the tree traversal.

[ExpandingMan]

This is only for StatefulBFS though, no?

[oscardssmith]

no, this is a stateful DFS where you just store a the parents in a list rather than a a cursor.

[ExpandingMan]

Ok, so are you considering abandoning the cursor approach? (I don't necessarily think this would be a bad idea because this problem runs very deep.) We need to plan this out a bit because it's going to be a bit of an ordeal to fix this so I want to get some idea of where we will land before putting the effort in. Frustratingly, the chances this can be fixed without breaking changes seems low since I will almost definitely have to change some of the type signatures, though we can discuss afterward whether we really need to treat these as breaking since they are extremely unlikely to break any real packages.

[oscardssmith]

My thought was to abandon cursors for tree traversals, but to leave them in the package. I also believe we can get not completely awful performance out of tree cursors in a non-breaking way (keep the type parameter but don't put stuff in it).

[ExpandingMan]

Ok, in that case I might wait a few days to see what you come up with before going in and ripping apart the cursors.

In theory it should be possible to still get good performance out of cursors for HasNodeType trees even without abusing the compiler like this, but as I indicated, it's going to be pretty tricky. Storing the iteration state of every node in the tree is a really tough problem. I'll have to look back at Keno's initial implementation but I don't think it was fundamentally different and I remember it being type unstable.

[ExpandingMan]

By the way, if you are able to implement a performant solution that doesn't suffer from this pathology and is type-stable for HasNodeType(), I don't think it would be entirely unreasonable to keep the cursors more-or-less as is. Note that it doesn't have to recompile for every tree topology, only for trees with different numbers of generations. To me the bigger issue is that it starts to become extremely expensive to compile for very deep trees.

So it's not necessarily a crazy implementation to have available, it's just totally unacceptable as a default.

I should be able to tweak it to get rid of this problem at least in the simplest cases such as TrivialCursor.

[oscardssmith]

https://github.com/JuliaCollections/AbstractTrees.jl/pull/118 is my first try for fixing it. It fixes the MWE by adding some type instability, but I think it's relatively good.