Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid unnecessary keyword arguments check in fallback rrule #308

Merged
merged 2 commits into from
Mar 24, 2021

Conversation

Keno
Copy link
Contributor

@Keno Keno commented Mar 7, 2021

The expansion rrule(::Any, ::Vararg{Any}; kwargs...) actually
generates two methods. One along the lines of:

(::typeof(Core.kwfunc(rrule))(kwargs, ::typeof(rrule), ::Any, ::Vararg{Any}) = nothing

and the other that just calls it:

rrule(a::Any, b::Vararg{Any}) = Core.kwfunc(rrule)(NamedTuple{}(), a, b...)

The compiler handles this fallback well, since it's used all over the place,
but the cost to infer it is non-zero. Of course, in the AD use case, this
fallback method is visited literally on every call, so saving a tiny amount
of inference/compile time actually leads to noticable improvements over
a whole AD problem.

src/rules.jl Outdated Show resolved Hide resolved
@Keno Keno force-pushed the kf/fallbackvarargs branch from fd3ada6 to 2174b67 Compare March 23, 2021 20:25
The expansion `rrule(::Any, ::Vararg{Any}; kwargs...)` actually
generates two methods. One along the lines of:

```
(::typeof(Core.kwfunc(rrule))(kwargs, ::typeof(rrule), ::Any, ::Vararg{Any}) = nothing
```

and the other that just calls it:

```
rrule(a::Any, b::Vararg{Any}) = Core.kwfunc(rrule)(NamedTuple{}(), a, b...)
```

The compiler handles this fallback well, since it's used all over the place,
but the cost to infer it is non-zero. Of course, in the AD use case, this
fallback method is visited literally on every call, so saving a tiny amount
of inference/compile time actually leads to noticable improvements over
a whole AD problem.
@Keno Keno force-pushed the kf/fallbackvarargs branch from 2174b67 to 2ead553 Compare March 23, 2021 20:49
@Keno
Copy link
Contributor Author

Keno commented Mar 23, 2021

@oxinabox should be good to go

@oxinabox
Copy link
Member

Will review and merge tomorrow

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants