Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
introduce
@nospecializeinfer macro to tell the compiler to avoid ex…
…cess inference
This commit introduces a new compiler annotation called `@nospecializeinfer`,
which allows us to request the compiler to avoid excessive inference.
\## `@nospecialize` mechanism
T discuss `@nospecializeinfer`, let's first understand the behavior of
`@nospecialize`.
Its docstring says that
> This is only a hint for the compiler to avoid excess code generation.
, and it works by suppressing dispatches with complex runtime
occurrences of the annotated arguments. This could be understood with
the example below:
```julia
julia> function call_func_itr(func, itr)
local r = 0
r += func(itr[1])
r += func(itr[2])
r += func(itr[3])
r
end;
julia> _isa = isa; # just for the sake of explanation, global variable to prevent inlining
julia> func_specialize(a) = _isa(a, Function);
julia> func_nospecialize(@nospecialize a) = _isa(a, Function);
julia> dispatchonly = Any[sin, muladd, nothing]; # untyped container can cause excessive runtime dispatch
julia> @code_typed call_func_itr(func_specialize, dispatchonly)
CodeInfo(
1 ─ %1 = π (0, Int64)
│ %2 = Base.arrayref(true, itr, 1)::Any
│ %3 = (func)(%2)::Any
│ %4 = (%1 + %3)::Any
│ %5 = Base.arrayref(true, itr, 2)::Any
│ %6 = (func)(%5)::Any
│ %7 = (%4 + %6)::Any
│ %8 = Base.arrayref(true, itr, 3)::Any
│ %9 = (func)(%8)::Any
│ %10 = (%7 + %9)::Any
└── return %10
) => Any
julia> @code_typed call_func_itr(func_nospecialize, dispatchonly)
CodeInfo(
1 ─ %1 = π (0, Int64)
│ %2 = Base.arrayref(true, itr, 1)::Any
│ %3 = invoke func(%2::Any)::Any
│ %4 = (%1 + %3)::Any
│ %5 = Base.arrayref(true, itr, 2)::Any
│ %6 = invoke func(%5::Any)::Any
│ %7 = (%4 + %6)::Any
│ %8 = Base.arrayref(true, itr, 3)::Any
│ %9 = invoke func(%8::Any)::Any
│ %10 = (%7 + %9)::Any
└── return %10
) => Any
```
The calls of `func_specialize` remain to be `:call` expression (so that
they are dispatched and compiled at runtime) while the calls of
`func_nospecialize` are resolved as `:invoke` expressions. This is
because `@nospecialize` requests the compiler to give up compiling
`func_nospecialize` with runtime argument types but with the declared
argument types, allowing `call_func_itr(func_nospecialize, dispatchonly)`
to avoid runtime dispatches and accompanying JIT compilations
(i.e. "excess code generation").
The difference is evident when checking `specializations`:
```julia
julia> call_func_itr(func_specialize, dispatchonly)
2
julia> length(Base.specializations(only(methods(func_specialize))))
3 # w/ runtime dispatch, multiple specializations
julia> call_func_itr(func_nospecialize, dispatchonly)
2
julia> length(Base.specializations(only(methods(func_nospecialize))))
1 # w/o runtime dispatch, the single specialization
```
The problem here is that it influences dispatch only, and does not
intervene into inference in anyway. So there is still a possibility of
"excess inference" when the compiler sees a considerable complexity of
argument types during inference:
```julia
julia> func_specialize(a) = _isa(a, Function); # redefine func to clear the specializations
julia> @Assert length(Base.specializations(only(methods(func_specialize)))) == 0;
julia> func_nospecialize(@nospecialize a) = _isa(a, Function); # redefine func to clear the specializations
julia> @Assert length(Base.specializations(only(methods(func_nospecialize)))) == 0;
julia> withinfernce = tuple(sin, muladd, "foo"); # typed container can cause excessive inference
julia> @time @code_typed call_func_itr(func_specialize, withinfernce);
0.000812 seconds (3.77 k allocations: 217.938 KiB, 94.34% compilation time)
julia> length(Base.specializations(only(methods(func_specialize))))
4 # multiple method instances inferred
julia> @time @code_typed call_func_itr(func_nospecialize, withinfernce);
0.000753 seconds (3.77 k allocations: 218.047 KiB, 92.42% compilation time)
julia> length(Base.specializations(only(methods(func_nospecialize))))
4 # multiple method instances inferred
```
The purpose of this PR is to implement a mechanism that allows us to
avoid excessive inference to reduce the compilation latency when
inference sees a considerable complexity of argument types.
\## Design
Here are some ideas to implement the functionality:
1. make `@nospecialize` block inference
2. add nospecializeinfer effect when `@nospecialize`d method is annotated as `@noinline`
3. implement as `@pure`-like boolean annotation to request nospecializeinfer effect on top of `@nospecialize`
4. implement as annotation that is orthogonal to `@nospecialize`
After trying 1 ~ 3., I decided to submit 3.
\### 1. make `@nospecialize` block inference
This is almost same as what Jameson has done at <vtjnash@8ab7b6b>.
It turned out that this approach performs very badly because some of
`@nospecialize`'d arguments still need inference to perform reasonably.
For example, it's obvious that the following definition of
`getindex(@nospecialize(t::Tuple), i::Int)` would perform very badly if
`@nospecialize` blocks inference, because of a lack of useful type
information for succeeding optimizations:
<https://github.com/JuliaLang/julia/blob/12d364e8249a07097a233ce7ea2886002459cc50/base/tuple.jl#L29-L30>
\### 2. add nospecializeinfer effect when `@nospecialize`d method is annotated as `@noinline`
The important observation is that we often use `@nospecialize` even when
we expect inference to forward type and constant information.
Adversely, we may be able to exploit the fact that we usually don't
expect inference to forward information to a callee when we annotate it
with `@noinline` (i.e. when adding `@noinline`, we're usually fine with
disabling inter-procedural optimizations other than resolving dispatch).
So the idea is to enable the inference suppression when `@nospecialize`'d
method is annotated as `@noinline` too.
It's a reasonable choice and can be efficiently implemented with #41922.
But it sounds a bit weird to me to associate no infer effect with
`@noinline`, and I also think there may be some cases we want to inline
a method while partly avoiding inference, e.g.:
```julia
\# the compiler will always infer with `f::Any`
@noinline function twof(@nospecialize(f), n) # this method body is very simple and should be eligible for inlining
if occursin('+', string(typeof(f).name.name::Symbol))
2 + n
elseif occursin('*', string(typeof(f).name.name::Symbol))
2n
else
zero(n)
end
end
```
\### 3. implement as `@pure`-like boolean annotation to request nospecializeinfer effect on top of `@nospecialize`
This is what this commit implements. It basically replaces the previous
`@noinline` flag with a newly-introduced annotation named `@nospecializeinfer`.
It is still associated with `@nospecialize` and it only has effect when
used together with `@nospecialize`, but now it is not associated to
`@noinline`, and it would help us reason about the behavior of `@nospecializeinfer`
and experiment its effect more safely:
```julia
\# the compiler will always infer with `f::Any`
Base.@nospecializeinfer function twof(@nospecialize(f), n) # the compiler may or not inline this method
if occursin('+', string(typeof(f).name.name::Symbol))
2 + n
elseif occursin('*', string(typeof(f).name.name::Symbol))
2n
else
zero(n)
end
end
```
\### 4. implement as annotation that is orthogonal to `@nospecialize`
Actually, we can have `@nospecialize` and `@nospecializeinfer` separately, and it
would allow us to configure compilation strategies in a more
fine-grained way.
```julia
function noinfspec(Base.@nospecializeinfer(f), @nospecialize(g))
...
end
```
I'm fine with this approach but at the same time I'm afraid to have too
many annotations that are related to some sort (I expect we will
annotate both `@nospecializeinfer` and `@nospecialize` in this scheme).
Co-authored-by: Mosè Giordano <giordano@users.noreply.github.com>
Co-authored-by: Tim Holy <tim.holy@gmail.com>- Loading branch information
3 people
committed
May 10, 2023
1 parent
b9b8b38
commit 63dab32