Trait dispatch causing allocations

I am simulating electronic circuits and I have chosen to adopt a modular model which works on objects of type Component. A Component can be a Resistor, Capacitor, or even a Serial/Parallel combination of Components. This is reflected in the following code block:

abstract type Component end

mutable struct Resistor <: Component end

mutable struct Capacitor <: Component end

struct Serial{P} <:Component 
    components::P
    Serial(v::Vararg{<:Component,N}) where N = new{typeof(v)}(v)
end

struct Parallel{P} <:Component 
    components::P
    Parallel(v::Vararg{<:Component,N}) where N = new{typeof(v)}(v)
end

I would like to apply a generic function (applyf) to a Component, but I would like to distinguish between Components that hold a composition of Components and a singular Component. To achieve this, I have adopted the `Holy Traits’ technique:

abstract type CompositionTrait end
struct Leaf <: CompositionTrait end
struct Composite{N} <: CompositionTrait end

CompositionTrait(::Type{<:Resistor}) = Leaf()
CompositionTrait(::Type{<:Capacitor}) = Leaf()
CompositionTrait(::Type{<:Serial{P}}) where P = Composite{length(P.parameters)}()
CompositionTrait(::Type{<:Parallel{P}}) where P  = Composite{length(P.parameters)}()

Great! I can now write different functions for a Component which has the trait Composite{N} and a Component which has the trait Leaf.

applyf(c::Component) = applyf(CompositionTrait(typeof(c)), c)

applyf(::Leaf, c::Component) = nothing

function applyf(::Composite{N},  c::Component)  where N
    for i = 1:N
        applyf(c.components[i])
    end
    nothing
end

This works as expected. Moreover, for simple compositions of Components, the code is very efficient with zero allocations. However, the following complicated (nested compositions) component causes allocations:

using BenchmarkTools
const N = 10
component = 
Serial(
    (Parallel(
        (Serial(Resistor(),Capacitor()) for _=1:N)...
    ) for i=1:N
    )...
)


>julia @btime applyf($component)
'316.843 ns (10 allocations: 1.72 KiB)'

When starting Julia with --track-allocation=user, it notifies me that the trait dispatch line applyf(c::Component), is causing the allocations. Why is this? The Components with trait Composite{N}, hold a tuple of Components so all type information should be resolved at compile-time.

Apologies in advance for the longest MWE you will have seen… Thanks!

1 Like

I think you are reaching the cases where run time dispatch is kicking in. This kind of stuff is discussed in this thread and other ones linked there, and I don’t think there is an ideal solution.

1 Like

I think these lines are your issue. You’re getting the length of the tuple from the type and using that to create the N parameter of the Composite instance. The problem is this is all happening at run time when you do length(P.parameters). Since this info is known at compile time, you could use @generated to make this happen at compile time.

@generated CompositionTrait(::Type{<:Serial{P}}) where P = :($(Composite{length(P.parameters)})())
@generated CompositionTrait(::Type{<:Parallel{P}}) where P  = :($(Composite{length(P.parameters)})())

But maybe more importantly, I’m not sure you actually need to do any of this. I think you can just drop the N parameter from Composite and use map to loop over the c.components, right?

Here’s how I might set it up:

abstract type Component end
mutable struct Resistor <: Component end
mutable struct Capacitor <: Component end

struct Serial{P} <: Component
    components::P
end
Serial(args...) = Serial(args)

struct Parallel{P} <: Component
    components::P
end
Parallel(args...) = Parallel(args)

abstract type CompositionTrait end
struct Leaf <: CompositionTrait end
struct Composite <: CompositionTrait end

CompositionTrait(::Type{<:Resistor}) = Leaf()
CompositionTrait(::Type{<:Capacitor}) = Leaf()
CompositionTrait(::Type{<:Serial}) = Composite()
CompositionTrait(::Type{<:Parallel}) = Composite()

applyf(c::Component) = applyf(CompositionTrait(typeof(c)), c)
applyf(::Leaf, c::Component) = nothing
applyf(::Composite,  c::Component) = map(applyf, c.components)
using BenchmarkTools
const N = 10
component = 
Serial(
    (Parallel(
        (Serial(Resistor(),Capacitor()) for _=1:N)...
    ) for i=1:N
    )...
)
julia> @btime applyf($component);
  0.001 ns (0 allocations: 0 bytes)
3 Likes

any reason for mutability?

@rveltz Not sure, I just copied that part from OP’s example.

Thanks for pointing this out, though I thought length of a Tuple would work at compile time?

Thank you providing a solution, though I’m not sure if it fully resolves the problem. Take a look at this (even more complex) component:

component = 
Parallel((
Serial(
    (Parallel(
        (Serial(Resistor(),Capacitor()) for _=1:N)...
    ) for i=1:N
    )...
) for _=1:N)...,(Resistor() for _=1:N)...)
julia> @btime applyf($component);

  15.041 μs (29 allocations: 316.38 KiB)

I originally used foreach (similar to map), in my solutions, but I found that it allocates. I believe the reason for this is because c.components holds a tuple of Components which is an abstract type. Therefore, although we know the type of each element in the tuple, the type is not the same for each element and therefore we can’t fix the type of the iterator? (maybe this is mumbo jumbo but I couldn’t figure out why else it would allocate).

I have got rid of a lot of the implementation to produce a MWE. Resistor in fact has many fields which can change over its lifetime.

Thanks for your reply. I’ve had a look at this thread, and you are right it is similar to the things discussed. However, one difference (and an important difference) is that rather than holding a vector of Components, I am holding a Tuple. This difference means that the compiler should be able to statically resolve the dispatch of each Component in the Tuple. I would have thought that the compiler is smart enough to understand the immutability of the tuple and therefore optimise the dispatch at compile time accordingly.

Yeah, at the end that depends on some heuristic parameters. The compiler may union-split if the number of types is small, but depending on the structure of the code and number of different types it may give up. Note that by using a tuple you could perhaps “force” that specialization, but that would be accompanied by specialization of that function for every different type, thus creating further compilation costs. I don’t know (really) if that is a good idea there.

1 Like