Can I avoid dynamic dispatch here ? (recursive parameter in struct )

There has been a while that I sense I might have some misconceptions about dynamic dispatch, and maybe this use case represents my doubts well..
So.. find myself in the unpleasant situation of having to use abstract containers.

Consider the following MWE:

abstract type MyAbstract end

struct Beta <: MyAbstract end

struct Alpha{T<:MyAbstract} <: MyAbstract
    a::Int
    allcompanion::Vector{T}
end

function dosomething(alpha)
    companion = first(alpha.allcompanion)
    return dosomethingelse(alpha, companion)
end

function dosomethingelse(alpha::Alpha, companion::Alpha)
    return alpha.a + companionjob(companion)
end

function companionjob(companion::Alpha)
    return companion.a
end

function companionjob(companion::Beta)
    return 1
end

alphabeta = Alpha(1, [Beta(), Beta(), Beta()])

alphaalpha1 = Alpha(1, Alpha[])
alphaalpha2 = Alpha(2, Alpha[])
alphaalpha3 = Alpha(3, Alpha[])

push!(alphaalpha1.allcompanion, alphaalpha2, alphaalpha3)
push!(alphaalpha2.allcompanion, alphaalpha1, alphaalpha3)
push!(alphaalpha3.allcompanion, alphaalpha1, alphaalpha2)

Running @code_warntype dosomething(alphaalpha1) obviously complains

Unfortunately, dosomething is in a hot loop, and profiling yields much dynamic dispatch happening there..

I have a couple of questions:

  1. I think it’s crystal clear for a human processing what should happen in dosometing(::Alpha{Alpha}), It’s clear that companionjob(::Alpha) will be executed, and since we don’t ask for the allcompanion abstract field of the companion, no surprises are expected.. Do you agree? If yes, why does the compiler have problems then?

  2. Does the compiler really have problems, and dynamic dispatch really happens here such that there is a performance hit, or could that be a false positive, since, as I said in 1., all generated code should be predicted?

  3. What can I do to lift any performance issues with dynamic dispatch ? I am willing to go all the way for internals and hacks..

I think you just need not to create abstract containers Alpha[] e.g. alphaalpha1 = Alpha(1, Alpha{Alpha}[]), with this:


julia> @code_warntype dosomething(alphaalpha1)
MethodInstance for dosomething(::Alpha{Alpha{Alpha}})
  from dosomething(alpha) @ Main REPL[38]:1
Arguments
  #self#::Core.Const(dosomething)
  alpha::Alpha{Alpha{Alpha}}
Locals
  companion::Alpha{Alpha}
Body::Int64
1 ─ %1 = Base.getproperty(alpha, :allcompanion)::Vector{Alpha{Alpha}}
│        (companion = Main.first(%1))
│   %3 = Main.dosomethingelse(alpha, companion)::Int64
└──      return %3

and if you need also Beta you can use Union{Alpha{Alpha}, Alpha{Beta}}[].

edit: though, if you have more nesting I guess you will still have problems, so this is not a solution, more that the MWE is not very good…my guess is that there is no simple solution if you have nested parameters because the number of types explode, an idea would be using instead something like

struct Alpha <: MyAbstract
    a::Int
    allcompanion::Vector{Union{Alpha, Beta}}
end
1 Like

There weren’t any. The compiler branches over your few methods and infers a concrete Int64. Just goes to show that an inferred abstract type isn’t a strict dead end.

julia> using JET, BenchmarkTools

julia> @report_opt dosomething(alphaalpha1)
No errors detected

julia> @btime dosomething($alphaalpha1)
  2.000 ns (0 allocations: 0 bytes)
3

On a more nested level though, Alpha is still an abstract parameter and there’s a Vector{Alpha} somewhere. But nesting concrete types all the way down in Alpha{Alpha{#=circular parameter=#}} is just Zeno’s paradox; at some point, a concrete Alpha{Beta} or the abstract Alpha needs to happen. Ideally we don’t design types like this at all.

2 Likes

Consider what the implication is when you put the infromation about the children’s type into the type domain: You are now essentially encoding the entire tree in the type domain because the each node has to know the exact type of its children all the way down.
Additionally, the way you did it here (by choosing a Vec for the children) means that all children need to have the same type, i.e. need to have the depth in the tree.

While there might be usecases where this behavior is desireable (maybe graphs modeling a computarion where you want maximal performance and thus essentially try to compile a function for each specific graph), it is likely that you could get away with a simpler approach:
Just have 1 type of node that can express all the cases you need and then use regular if statements to branch on the node’s kind during runtime.

If you happen to have many different node types, then you could think about using LightSumTypes.jl sum types to separate the node types and generate the if statements for you. Although I am not sure from the top of my head whether this library works with recursive definitions.

1 Like

Tortar has made a better variant of his package, called WrappedUnions.jl.

2 Likes

Thanks everybody for chiming in. To elaborate on the use case, the allcompanion are basically handlers though which I can communicate to other Alpha objects. E.g., HTTP handlers. But for simulations, I would like to store the Alpha object as a handler as well..


@Benny That’s quite hopeful, actually. And this is one of my major questions regarding when dynamic dispatch happens. E.g., imagine I have this code

function entryfun()
    unstableobj = stablefun() # return a stable type (imagine Float64)
    stableobj = unstablefun() # return an unstable type
    isthisdynamicdispatchedfun(stableobj, unstableobj)
end

function isthisdynamicdispatchedfun(stableobj, unstableobj)
    return stableobj * 15
end

And similarly for structs

struct StableUnstableStruct
    stablefield::Float64
    unbstablefield::Any
end

function entryfun()
    sus = StableUnstableStruct(stablefun(), unstablefun())
    isthisdynamicdispatchedfun(sus)
end

function isthisdynamicdispatchedfun(sus::StableUnstableStruct)
    return sus.stablefield * 15
end

Then I guess, that even if the all arguments for isthisdynamicdispatchedfun are only resolved during runtime, the compiler knows that these are not needed inside the function, so it can already JIT compile the correct version without needing dynamic dispatch. Right ? If JET.@report_opt is the definitive tool to check for dynamic dispatch then that’s true..


@abraemer I think I like this approach because that will hold for any nested allcompanion access.
So the previous example will be converted to

abstract type MyAbstract end

struct Beta <: MyAbstract end

struct Alpha <: MyAbstract # <--- this changed: no more parameter
    a::Int
    allcompanion::Vector{MyAbstract} # <---- this changed: to be generic
end

function dosomething(alpha)
    companion = first(alpha.allcompanion)
    if companion isa Alpha
        return dosomethingelse(alpha, companion)
    elseif companion isa Beta
        return dosomethingelse(alpha, companion)
    else
        return dosomethingelse(alpha, companion)
    end
end

function dosomethingelse(alpha::Alpha, companion::MyAbstract)
    if companion isa Alpha
        companionresult = companionjob(companion)
    elseif companion isa Beta
        companionresult = companionjob(companion)
    else
        companionresult = companionjob(companion)
    end
    return alpha.a + companionresult
end

function companionjob(companion::Alpha)
    return companion.a
end

function companionjob(companion::Beta)
    return 1
end

alphabeta = Alpha(1, [Beta(), Beta(), Beta()])

alphaalpha1 = Alpha(1, Alpha[])
alphaalpha2 = Alpha(2, Alpha[])
alphaalpha3 = Alpha(3, Alpha[])

push!(alphaalpha1.allcompanion, alphaalpha2, alphaalpha3)
push!(alphaalpha2.allcompanion, alphaalpha1, alphaalpha3)
push!(alphaalpha3.allcompanion, alphaalpha1, alphaalpha2)

Okey. Obviously, JET now complains for the generic case (else):

However, I am fine with this, since it is essentially dead code and will never be executed.
I guess there can never be performance hit due to dynamic dispatch or similar for unreachable code, right ? :eyes: :eyes: :eyes:
I just created this generic case to eliminate the Union{Nothing, ...} being propagated everywhere..

I mean, I can even make the following macro

"""
Must use `companion` in the `expr` as the variable to use in the if statements.
"""
macro companionifexpand(expr)
    return quote
        if $(esc(:(companion isa Alpha)))
            $(esc(expr))
        elseif $(esc(:(companion isa Beta)))
            $(esc(expr))
        else
            $(esc(expr))
        end
    end
end

And then all the function become as easy to write as:

function dosomething(alpha)
    companion = first(alpha.allcompanion)
    @companionifexpand dosomethingelse(alpha, companion)
end

function dosomethingelse(alpha::Alpha, companion::MyAbstract)
    companionresult = @companionifexpand companionjob(companion)
    return alpha.a + companionresult
end

@abraemer, is that what you had in mind or should I consider something else that, e.g., SumTypes.jl, LightSumTypes.jl , or WrappedUnions.jl implement?

Yes that is what I had in mind. The very light-weight ( < 100 loc) package WrappedUnions.jl essentially provides a generic case for this and also provides a direct pendant to your companionifexpand macro. I encourage you to take a look at that :slight_smile: