Type instability in recursion

Hi, I am implementing genetic programming (GP) in julia. GP needs to evaluate mathematical expressions represented by tree structures, and this is the bottleneck in my implementation.
The following code can evaluate tree structures, but the type is not stable.

abstract type Primitive end
abstract type Terminal <: Primitive end

struct PrimitiveFunction <: Primitive
    func::Function
    arity::Int
end

struct Variable <: Terminal
    index::Int
end

struct Const{T <: Real} <: Terminal
    value::T
end

struct GPNode{T <: Primitive}
    primitive::T
    children::Vector{GPNode}
end
GPNode(prim::Terminal) = GPNode(prim, GPNode[])

function evaluate(node::GPNode{PrimitiveFunction}, vars)
    args = evaluate.(node.children, Ref(vars))
    node.primitive.func(args...)
end
evaluate(node::GPNode{Const{T}}, vars) where T <: Real =  node.primitive.value
evaluate(node::GPNode{Variable}, vars) = vars[node.primitive.index]


x = GPNode(Variable(1))
c = GPNode(Const(1.0))
root = GPNode(PrimitiveFunction(+, 2), [x, c])
vars = [2.0]
@code_warntype evaluate(root, vars)

And this is the result of @code_warntype .

Variables
  #self#::Core.Compiler.Const(evaluate, false)
  node::GPNode{PrimitiveFunction}
  vars::Array{Float64,1}
  args::Any

Body::Any
1 ─ %1 = Base.getproperty(node, :children)::Array{GPNode,1}
│   %2 = Main.Ref(vars)::Base.RefValue{Array{Float64,1}}
│   %3 = Base.broadcasted(Main.evaluate, %1, %2)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1},Nothing,typeof(evaluate),Tuple{Array{GPNode,1},Base.RefValue{Array{Float64,1}}}}
│        (args = Base.materialize(%3))
│   %5 = Base.getproperty(node, :primitive)::PrimitiveFunction
│   %6 = Base.getproperty(%5, :func)::Function
│   %7 = Core._apply_iterate(Base.iterate, %6, args)::Any
└──      return %7

I think this instability is slowing down the execution. How can I get rid of this instability?

You can’t get rid of it if you keep PrimitiveFunction.func as Function because that’s an abstract type. Inference can’t predict what that function will return. You could use a type parameter to specialise on this, but that will probably kill your compile time:

struct PrimitiveFunction{T} <: Primitive
    func::T
    arity::Int
end

Also, your GPNode.children would have to stay as a very generic container, since the type parameter would specialise on functions and thus quickly devolve into Any as a container type.

1 Like

Why would using a type parameter kill compile time? I don’t think that’s a major compile time bottleneck.

Not one specific compile time, but overall compile time of the program would increase, since each function has its own type:

julia> typeof(sin) == typeof(cos)
false

So specializing on function type would mean having to compile the functions using the type for every different function that is wrapped, which would lead to an overall increase in compile time.

(Unless @nospecialize is used, in which case the type parameter is probably no longer useful…)

I assumed that compiling the boxed unstable functions will take a lot longer than compiling the type-stable methods - it’s so much more native code in the end.

If you are only using a few types the type-stable methods should be faster overall?

Since the code is calling itself recursively here, there might be a lot of compilation happening just to compile the evaluate tree. In the current implementation, evaluate(::GPNode{PrimitiveType}) is only compiled once, since PrimitiveType is not specialized. That’s what would increase the compile time.

Yes, presumably. Note that at some point inference will give up even with type stable methods, since the broadcast has to unify the return types somehow. It will rather soon converge to Any, which will introduce a lot of dynamic lookups when calling func anyway…

1 Like

Oh yeah you should never use broadcast for something like this - I missed that part of the code. Unless those vectors are huge this should be all tuples and map. Otherwise evaluate can’t be type-stable.

But anecdotally lot of my packages use this kind of recursion on tuples basically everywhere, and the compile times are fine.

If vectors are large there are wyas for reducing the cost of dynamic method dispatch - there are some recent comments by tim holy about this somewhere.

Are your packages specializing on functions? I hope I was clear about this applying if OP wanted to make their code type stable by using a type parameter for the function type in PrimitiveFunction :slight_smile:

On tuples the story is similar, since each combination of element types leads to a unique tuple type. If you limit your code to a small number of types and tuple lengths, it will be fine, but since OP is (seemingly) building a very generic framework, type explosion can easily occur.

1 Like

In this case it depends on how many functions he is using. I would assume its fine unless it’s a LOT. Like 100, sure - that might need a different approach. But 10 is fine.

The question here is the relative costs of type instability and of a little compilation. The question here is also about removing type instability - so I’m assuming that’s the main problem in this specific context.

TLDR;

@kogad - Focusing only on the problem of type stability, you need to use concrete types with type parameters, and map over tuples rather than broadcast over vectors. Mixed type vectors like your children field are not type-stable.

There may be consequences for compilation of this approach, that may or may not matter in your context.

2 Likes

Thanks for your response!
I’m going to use a type parameter and tuples.

Worth also noting that if your Functions have predictable signatures, you can use FunctionWrappers.jl. This will create some runtime cost, but reduce compile time. Depends on what’s more important.