Update kwargs fallback rules after RuleConfig rewrite

Fixes #368

src/rule_definition_tools.jl

@@ -342,14 +342,16 @@ end
 
 function _nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke)
     @gensym kwargs
-    # `::Any` instead of `_`: https://github.com/JuliaLang/julia/issues/32727
     return @strip_linenos quote
-        function ChainRulesCore.frule(
-            @nospecialize(::Any), $(map(esc, primal_sig_parts)...); $(esc(kwargs))...
-        )
+        # Manually defined kw version to save compiler work. See explanation in rules.jl
+        function (::Core.kwftype(typeof(ChainRulesCore.frule)))(@nospecialize($kwargs::Any),
+                frule::typeof(ChainRulesCore.frule), @nospecialize(::Any), $(map(esc, primal_sig_parts)...))
+            return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), NoTangent())
+        end
+        function ChainRulesCore.frule(@nospecialize(::Any), $(map(esc, primal_sig_parts)...))
             $(__source__)
             # Julia functions always only have 1 output, so return a single NoTangent()
-            return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), NoTangent())
+            return ($(esc(primal_invoke)), NoTangent())
         end
     end
 end

src/rules.jl

@@ -58,10 +58,25 @@ will be hit as a fallback. This is the case for most rules.
 
 See also: [`rrule`](@ref), [`@scalar_rule`](@ref), [`RuleConfig`](@ref)
 """
-frule(::Any, ::Any, ::Vararg{Any}; kwargs...) = nothing
+frule(::Any, ::Any, ::Vararg{Any}) = nothing
 
 # if no config is present then fallback to config-less rules
-frule(::RuleConfig, ȧrgs, f, args...; kwargs...) = frule(ȧrgs, f, args...; kwargs...)
+frule(::RuleConfig, args...) = frule(args...)
+
+# Manual fallback for keyword arguments. Usually this would be generated by
+#
+#   frule(::Any, ::Vararg{Any}; kwargs...) = nothing
+#
+# However - the fallback method is so hot that we want to avoid any extra code
+# that would be required to have the automatically generated method package up
+# the keyword arguments (which the optimizer will throw away, but the compiler
+# still has to manually analyze). Manually declare this method with an
+# explicitly empty body to save the compiler that work.
+const frule_kwfunc = Core.kwftype(typeof(frule)).instance
+(::typeof(frule_kwfunc))(::Any, ::typeof(frule), ::Any, ::Vararg{Any}) = nothing
+(::typeof(frule_kwfunc))(kws::Any, ::typeof(frule), ::RuleConfig, args...) =
+    (frule_kwfunc)(kws, frule, args...)
+
 
 """
     rrule([::RuleConfig,] f, x...)
@@ -116,18 +131,10 @@ See also: [`frule`](@ref), [`@scalar_rule`](@ref), [`RuleConfig`](@ref)
 rrule(::Any, ::Vararg{Any}) = nothing
 
 # if no config is present then fallback to config-less rules
-rrule(::RuleConfig, f, args...; kwargs...) = rrule(f, args...; kwargs...)
-# TODO do we need to do something for kwargs special here for performance?
-# See: https://github.com/JuliaDiff/ChainRulesCore.jl/issues/368
-
-# Manual fallback for keyword arguments. Usually this would be generated by
-#
-#   rrule(::Any, ::Vararg{Any}; kwargs...) = nothing
-#
-# However - the fallback method is so hot that we want to avoid any extra code
-# that would be required to have the automatically generated method package up
-# the keyword arguments (which the optimizer will throw away, but the compiler
-# still has to manually analyze). Manually declare this method with an
-# explicitly empty body to save the compiler that work.
+rrule(::RuleConfig, f, args...) = rrule(f, args...)
 
-(::Core.kwftype(typeof(rrule)))(::Any, ::Any, ::Vararg{Any}) = nothing
+# Manual fallback for keyword arguments. See above
+const rrule_kwfunc = Core.kwftype(typeof(rrule)).instance
+(::typeof(rrule_kwfunc))(::Any, ::typeof(rrule), ::Any, ::Vararg{Any}) = nothing
+(::typeof(rrule_kwfunc))(kws::Any, ::typeof(rrule), ::RuleConfig, args...) =
+    (rrule_kwfunc)(kws, rrule, args...)