Need help on understanding type inference / performance

Hi there,

I am the author of ExtensibleEffects.jl and several people approached me that it would be great to improve performance (it is currently super slow).

I made a first attempt on inspecting and found the following specific case to be one crucial component

struct ChainedFunctions{Fs}
    functions::Fs
    ChainedFunctions(functions...) = new{typeof(functions)}(functions)
end

instantiate(T) = T()

function instantiate_wrapper1(type)
    continuation = ChainedFunctions(() -> instantiate(type))
    first_func = Base.first(continuation.functions)
    first_func()
end

function instantiate_wrapper2(::Type{type}) where type
    continuation = ChainedFunctions(() -> instantiate(type))
    first_func = Base.first(continuation.functions)
    first_func()
end

using Test
@inferred instantiate_wrapper1(Vector) 
# ERROR: return type Vector{Any} does not match inferred return type Any
@inferred instantiate_wrapper2(Vector)
# Any[]

This is run on julia 1.7.1. It is quite a pity, as to the best of my understanding it should never make a type-inference difference in normal functions whether you use type::Type or ::Type{type} where type.

Furthermore, this distinction makes it impossible to write generic code, as out of a sudden, passing a type like Vector is something completely different than passing a concrete struct instance.

Motivation: In ExtensibleEffects Vector and others are effect-handlers. While you can have a generic effect handler implementation, like indeed for Vector, there may be the need for extra information in order to run the handler. The latter requires a concrete handler struct which captures those extra information. Both kinds of effect handlers are very intuitive and I would not like fallback to something like Val{Vector}() to support Vector - it would complicate the entire interface at many places.


Any help is highly appreciated. If someone knows, why this is going on, or how a general workaround could look like, or whether this may be unintended and worth a feature request, or even a bug report, everything is very much welcome.

1 Like

BAM! I just now imagined the workaround:

function instantiate_wrapper3(::Type{type}) where type
    instantiate_wrapper1(type)
end
@inferred instantiate_wrapper3(Vector)
# Any[]

Very surprising, that you have to specify such a fallback for type-inference to work correctly.
For me it looks like this wrapper shouldn’t be necessary.

EDIT: this workaround apparently fails very quickly. See my below reply

It’s presumably related to this:
https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing

I.e. Julia intentionally does not specialize instantiate_wrapper1 because the input type is not used directly in the function. However, I do agree that this behaviour is a bit of a footgun, or at least can be perplexing.

3 Likes

thank you for the reference.

can you explain, why the wrapper is enough to make it work? The inner function should still not specialize, or should it?

Unfortunately, this workaround seems to be a special case and already does not work any longer for two arguments

singleton_array(T::Type{<:AbstractArray}, a) = convert(T, [a]) 

struct ChainedFunctions{Fs}
    functions::Fs
    ChainedFunctions(functions...) = new{typeof(functions)}(functions)
end

function singleton_wrapper1(type, value)
    continuation = ChainedFunctions(x -> singleton_array(type, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper2(type::T, value) where T
    continuation = ChainedFunctions(x -> singleton_array(type, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper3(::Type{type}, value) where type
    continuation = ChainedFunctions(x -> singleton_array(type, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper1_wrapper1(::Type{type}, value) where type
    singleton_wrapper1(type, value)
end
function singleton_wrapper1_wrapper1(type, value)
    singleton_wrapper1(type, value)
end

@inferred singleton_wrapper1(Vector, 1)
# ERROR: return type Vector{Int64} does not match inferred return type Any
@inferred singleton_wrapper2(Vector, 1)
# ERROR: return type Vector{Int64} does not match inferred return type Any
@inferred singleton_wrapper3(Vector, 1)
# [1]

@inferred singleton_wrapper1_wrapper1(Vector, 1)
# ERROR: return type Vector{Int64} does not match inferred return type Any

How can I trigger the specialization on Type{T} while still keeping the generic function which works on other inputs than Type?

Can you explain a little better what your code intend to do? It seems to me that you have a struct for which your type parameter is often a Vector containing multiple distinct functions, and in your you retrieve one (the first?) of these function and apply to another input, how exactly you expect Julia to infer anything? It seems magical to me that Julia is able to infer anything at all, and if it does is probably because constant propagation, no?

I’m not sure I understand what you are doing, but the following seems to work for me on Julia 1.9.0

singleton_array(T::Type{<:AbstractArray}, a) = convert(T, [a]) 

struct ChainedFunctions{Fs}
    functions::Fs
    ChainedFunctions(functions...) = new{typeof(functions)}(functions)
end

function singleton_wrapper1(type, value)
    continuation = ChainedFunctions(x -> singleton_array(type, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper2(type::T, value) where T
    continuation = ChainedFunctions(x -> singleton_array(T, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper3(::Type{T}, value) where T
    continuation = ChainedFunctions(x -> singleton_array(T, x))
    first_func = Base.first(continuation.functions)
    first_func(value)
end

function singleton_wrapper_wrapper1(::Type{T}, value) where T
    singleton_wrapper1(T, value)
end
function singleton_wrapper_wrapper1(type, value)
    singleton_wrapper1(type, value)
end

function singleton_wrapper_wrapper3(::Type{T}, value) where T
    singleton_wrapper3(T, value)
end
function singleton_wrapper_wrapper3(type, value)
    singleton_wrapper3(type, value)
end

using Test
# @inferred singleton_wrapper1(Vector, 1) # ERROR: return type Vector{Int64} does not match inferred return type Any
# @inferred singleton_wrapper2(Vector, 1) #ERROR: MethodError: no method matching singleton_array(::Type{UnionAll}, ::Int64)
@inferred singleton_wrapper3(Vector, 1) #works
# @inferred singleton_wrapper_wrapper1(Vector, 1) # ERROR: return type Vector{Int64} does not match inferred return type Any
@inferred singleton_wrapper_wrapper3(Vector, 1)

Julia can infer everything if it specializes the function also over types.

The type parameter does not contain a Vector but a Tuple which indeed captures all relevant type information. The only type-information which is lost is that of the handler, and that gets lost because of missing specialization apparently.

Hi @goerch thank you for testing on 1.9

I get that singleton_wrapper_wrapper3 works on 1.9, which sounds great. Would be nice to have it for 1.6 so that I could support julia’s longterm release.
Thank you for your help.

Thank you all for your help. It seems the reason is understood: Missing type-specialization is the reason.

I still struggle with how to enable this and created a more concrete, hopefully also clearer, follow up question: How to mark a function argument for specialization on types while still staying as a generic argument?

1 Like

Is it really necessary to work on Type{...}-signatures? It can sometimes confuse inference very much especially and thus should generally be avoided if possible as described in the performance tip @jakobnissen posted above. A note here would be that additional type parameter causes specialization by giving up inference and resorting to dynamic dispatch.

1 Like

I wrote my code without the Type{...} signature which unfortunately results in very slow performance.

as @jakobnissen pointed out, this is due to Julia compiler requiring the Type{...} signature. It is not me who makes this necessary, it is the current Julia compiler.

Ah, ok, it did not become clear for me anywhere that it was a tuple. Yes, if it is a tuple then the type of the elements is not simplified to a common denominator and this can work.

Thank you all for your input

I filed a proper report at Bug Report: Type argument not specializing for closures · Issue #45257 · JuliaLang/julia · GitHub