Unexpected non-specialization for `DataType` arguments

[mikmoore]

See https://discourse.julialang.org/t/extra-allocation-with-t-datatype/85930.

According to https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing,

As a heuristic, Julia avoids automatically specializing on argument type parameters in three specific cases: Type, Function, and Vararg. Julia will always specialize when the argument is used within the method, but not if the argument is just passed through to another function.

However, the quoted text appears to be misleading.

julia> versioninfo()
Julia Version 1.8.0
Commit 5544a0fab76 (2022-08-17 13:38 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-9850H CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, skylake)
  Threads: 1 on 12 virtual cores

julia> foo(T,x) = exp(T(x))
foo (generic function with 1 method)

julia> foo(x) = exp(Float64(x))
foo (generic function with 2 methods)

julia> bar(f,x) = exp(f(x))
bar (generic function with 1 method)

julia> bar(x) = exp(sin(x))
bar (generic function with 2 methods)

julia> baz(x::Vararg) = +(x...)
baz (generic function with 1 method)

julia> using BenchmarkTools

julia> @btime foo($Float64,$10)
  88.297 ns (1 allocation: 16 bytes) # fails to specialize
22026.465794806718

julia> @btime foo($10)
  11.411 ns (0 allocations: 0 bytes) # ok
22026.465794806718

julia> @btime bar($sin,$10)
  26.432 ns (0 allocations: 0 bytes) # ok
0.5804096620472413

julia> @btime bar($10)
  25.803 ns (0 allocations: 0 bytes) # ok
0.5804096620472413

julia> @btime baz($1,$2,$3e0)
  2.200 ns (0 allocations: 0 bytes)
6.0

So it would appear that Function and Vararg arguments obey the documented behavior but DataType arguments are reluctant to specialize even when directly used within the function. Is this a consequence of typeof(Float64) === DataType whereas for functions we only get a typeof(sin) <: Function relationship?

It would seem that either a fix to the behavior or a clarification to the documentation is warranted.

[lucifer1004]

According to Benny's experiment, T::Type and T::Any are specialized but there are still extra allocations.

So there might be two separate issues. One is that T::DataType is not specialized, and the other is that using types as arguments causes unnecessary allocations.

[vtjnash]

This is just an artifact of incorrect benchmarking usage. Usually what you want to use with $ is some Ref{T}, since otherwise you are still leaving it open to doing significant constant-prop.

julia> @btime bar($(Ref{Function}(sin))[],$10)
  41.201 ns (1 allocation: 16 bytes)
0.5804096620472413

julia> @btime bar($(Ref{Any}(Float64))[],$10)
  113.490 ns (1 allocation: 16 bytes)
22026.465794806718

julia> @btime bar($(Ref{Type{Float64}}(Float64))[],$10)
  10.993 ns (0 allocations: 0 bytes)
22026.465794806718

julia> @btime bar($(Ref{DataType}(Float64))[],$10)
  114.558 ns (1 allocation: 16 bytes)
22026.465794806718