`@nospecialize` for variadic and keyword arguments

General Usage Performance
, , ,
1

I have a function pt:

function pt(
    args::Union{AbstractFuture,PartitionType,PartitionTypeComposition,Vector}...;
    kwargs...,
)

that has a whole bunch of specializations: (determined using SnoopCompile):

var"#pt#214"(kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), args::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}...) in Banyan at /home/calebwin/Projects/banyan-julia/Banyan/src/annotation.jl:273 (33 specializations)
julia> collect_for(mref[], banyandf_inf)
33-element Vector{SnoopCompileCore.InferenceTimingNode}:
 InferenceTimingNode: 0.001690/0.245240 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, _A, Tuple{Symbol, Symbol}, NamedTuple{names, T}} where {_A, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), ::AbstractFuture, ::PartitionType) with 9 direct children
 InferenceTimingNode: 0.007639/0.039909 on Banyan.var"#pt#214"(::Base.Pairs, ::typeof(pt), ::AbstractFuture, ::Banyan.PartitionTypeComposition) with 17 direct children
 InferenceTimingNode: 0.000487/0.039676 on Banyan.var"#pt#214"(::Base.Pairs, ::typeof(pt), ::PartitionType, ::Banyan.PartitionTypeComposition) with 5 direct children
 InferenceTimingNode: 0.000468/0.019484 on Banyan.var"#pt#214"(::Base.Pairs, ::typeof(pt), ::Banyan.PartitionTypeComposition, ::Banyan.PartitionTypeComposition) with 6 direct children
 InferenceTimingNode: 0.007155/0.028872 on Banyan.var"#pt#214"(::Base.Pairs, ::typeof(pt), ::Vector, ::Banyan.PartitionTypeComposition) with 4 direct children
 InferenceTimingNode: 0.006646/0.009635 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), ::AbstractFuture, ::Banyan.PartitionTypeComposition) with 11 direct children
 InferenceTimingNode: 0.000329/0.000329 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), ::PartitionType, ::Banyan.PartitionTypeComposition) with 0 direct children
 InferenceTimingNode: 0.000305/0.000305 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), ::Banyan.PartitionTypeComposition, ::Banyan.PartitionTypeComposition) with 0 direct children
 InferenceTimingNode: 0.006276/0.006276 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(pt), ::Vector, ::Banyan.PartitionTypeComposition) with 0 direct children
 InferenceTimingNode: 0.001258/0.021914 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::AbstractFuture, ::Vector{PartitionType}) with 4 direct children
 InferenceTimingNode: 0.000505/0.033146 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::PartitionType, ::Vector{PartitionType}) with 4 direct children
 InferenceTimingNode: 0.000522/0.025492 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Banyan.PartitionTypeComposition, ::Vector{PartitionType}) with 4 direct children
 InferenceTimingNode: 0.001459/0.024593 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Vector, ::Vector{PartitionType}) with 4 direct children
 InferenceTimingNode: 0.001763/0.028671 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Vector{PartitionType}) with 9 direct children
 InferenceTimingNode: 0.002113/0.145384 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::DataFrame, ::PartitionType) with 8 direct children
 InferenceTimingNode: 0.005611/0.106553 on Banyan.var"#pt#214"(::Base.Pairs, ::typeof(pt), ::Future, ::Banyan.PartitionTypeComposition) with 11 direct children
 InferenceTimingNode: 0.002113/0.027297 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::DataFrame, ::Vector{PartitionType}) with 5 direct children
 InferenceTimingNode: 0.001766/0.002103 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::DataFrame, ::PartitionType) with 3 direct children
 InferenceTimingNode: 0.010461/0.362331 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::DataFrame, ::AbstractFuture) with 30 direct children
 InferenceTimingNode: 0.005548/0.036449 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::DataFrame, ::Banyan.PartitionTypeComposition) with 7 direct children
 InferenceTimingNode: 0.001690/0.031525 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::DataFrame, ::Vector) with 5 direct children
 InferenceTimingNode: 0.008176/0.009496 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::DataFrame, ::AbstractFuture) with 10 direct children
 InferenceTimingNode: 0.005053/0.005053 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::DataFrame, ::Banyan.PartitionTypeComposition) with 0 direct children
 InferenceTimingNode: 0.001352/0.001352 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::DataFrame, ::Vector) with 0 direct children
 InferenceTimingNode: 0.009548/0.044886 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::DataFrame, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 19 direct children
 InferenceTimingNode: 0.009624/0.044412 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, String}}}, ::typeof(pt), ::Future, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 19 direct children
 InferenceTimingNode: 0.007820/0.016963 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::DataFrame, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 9 direct children
 InferenceTimingNode: 0.007758/0.016644 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Future, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 8 direct children
 InferenceTimingNode: 0.009723/0.020603 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::DataFrame, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 16 direct children
 InferenceTimingNode: 0.009555/0.020691 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol}, NamedTuple{(:match, :on), Tuple{DataFrame, Vector{String}}}}, ::typeof(pt), ::Future, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 16 direct children
 InferenceTimingNode: 0.007843/0.040927 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::GroupedDataFrame, ::Vararg{Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}}) with 11 direct children
 InferenceTimingNode: 0.002056/0.031967 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Union{AbstractFuture, PartitionType, Banyan.PartitionTypeComposition, Vector}, ::Vector{PartitionType}) with 9 direct children
 InferenceTimingNode: 0.002337/0.008347 on Banyan.var"#pt#214"(::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(pt), ::Future, ::Vector{PartitionType}) with 9 direct children

Should I be using @nospecialize in some way on these args... and kwargs...?

2

If you want it specialized for each combination of inputs, no, you should leave it like it is. But in many circumstances you can reduce compile latency with no worse runtime performance (indeed, sometimes better runtime performance) with @nospecialize. As the docstring for @nospecialize explains, you can just wrap each argument with it, or even easier just insert it on the first line of the body:

function foo(args...; kwargs...)
    @nospecialize      # without arguments, this nospecializes all arguments
    # body
end

Keep in mind that @nospecialize doesn’t prevent type-inference in cases where the caller knows the type more specifically. If you do need to block inference, you can wrap caller arguments with Base.inferencebarrier(x):

foo(Base.inferencebarrier(x); kwargs...)

Base.inferencebarrier(x) ensures that the return value is inferred as Any and effectively hides the true type of x from inference.

3

Thank you, Tim! Really appreciate your help with this and apologies if I’m asking stuff that is documented already. I’ve been reading many of the helpful resources you’ve contributed and using SnoopCompile as well and am very grateful for your contributions to this community. :heart:

Is it possible to use Base.inferencebarrier on kwargs? It looks like @nospecialize is still causing specialization for a function:

function DataFrames.groupby(df::DataFrame, cols::Any; kwargs...)::GroupedDataFrame
    # We will simply pass the `cols` and `kwargs` into `Future` constructor
    # so specialization isn't really needed
    @nospecialize
    # bunch of code here...
end

var"#groupby#397"(kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, ::typeof(groupby), df::DataFrame, cols) in BanyanDataFrames at /home/calebwin/Projects/banyan-julia/BanyanDataFrames/src/gdf.jl:36 (5 specializations)
julia> collect_for(mref[], banyandf_inf)
5-element Vector{SnoopCompileCore.InferenceTimingNode}:
 InferenceTimingNode: 0.044587/0.710611 on BanyanDataFrames.var"#groupby#397"(Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}()::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, groupby::typeof(groupby), ::DataFrame, ::Symbol) with 18 direct children
 InferenceTimingNode: 0.043023/0.523912 on BanyanDataFrames.var"#groupby#397"(::Base.Pairs, groupby::typeof(groupby), ::DataFrame, ::Any) with 3 direct children
 InferenceTimingNode: 0.048939/0.115327 on BanyanDataFrames.var"#groupby#397"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, groupby::typeof(groupby), ::DataFrame, ::Any) with 23 direct children
 InferenceTimingNode: 0.043784/0.043784 on BanyanDataFrames.var"#groupby#397"(Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}()::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, groupby::typeof(groupby), ::DataFrame, ::Any) with 0 direct children
 InferenceTimingNode: 0.061140/0.081147 on BanyanDataFrames.var"#groupby#397"(::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}, groupby::typeof(groupby), ::DataFrame, ::Any) with 24 direct children

I have a similar question for another function with a bunch of default argument values where @nospecialize itself didn’t help:

function configure(; user_id=nothing, api_key=nothing, ec2_key_pair_name=nothing, banyanconfig_path=nothing)
    @nospecialize
    # bunch of code here...
end

var"#configure#9"(user_id, api_key, ec2_key_pair_name, banyanconfig_path, ::typeof(configure)) in Banyan at /home/calebwin/Projects/banyan-julia/Banyan/src/utils.jl:129 (2 specializations)
julia> collect_for(mref[], banyandf_inf)
2-element Vector{SnoopCompileCore.InferenceTimingNode}:
 InferenceTimingNode: 0.005174/1.599911 on Banyan.var"#configure#9"(nothing::Nothing, nothing::Nothing, nothing::Nothing, nothing::Nothing, configure::typeof(configure)) with 9 direct children
 InferenceTimingNode: 0.005149/0.246636 on Banyan.var"#configure#9"(::Any, ::Any, ::Any, ::Any, configure::typeof(configure)) with 2 direct children

My program is currently dominated by inference time in the tens of seconds so I’m going to focus on adding concrete type annotations and improving inferrability of code right now. But I wanted to better understand how to despecialize functions where runtime type checking is okay.

4

Great questions! First, g(; @nospecialize(kwargs...)) may help for calls like f(; kwargs...) = g(; Base.inferencebarrier(kwargs)...) but may not for g(; color=:green). That’s because the compiler can see the specific keyword arguments in the second call.

Second, for certain types of keyword arguments it’s possible that constant-propagation is playing a role (check for yellow bars in ProfileView.view(flamegraph(tinf)); you could consider using

Base.@constprop :none g(; color=:blue, n=1) = "$color"^n

@constprop is in Compat so you can use it on versions of Julia before 1.7 (but it won’t have any effect).

Finally, starting with 1.8.0-beta2 there will be significant changes to precompilation, and you may be able to get rid of your inference just by precompilation if the types are predictable.