Fix inlining avoidance (#759)

* Failing test

* Failing tests in contexts

* Refactor primitive system

* Improve contexts testing

* Update low_level_maths not-a-primitive tests

* Update tools_for_rules is_primitive calls

* Formatting

* Fix test

* Fix doctests

* Expand on maybe_primitive semantics

* Hopefully improve performance

* Remove additional complexity and fix docs build

* Fix stuff

* Formatting

* Bump version

* Improve documentation:

* Reference to is_primitive replaced with _is_primitive

* Uncomment precompilation help

* Improve docstring

* Improve comment at use of maybe_primitive

* Update src/interpreter/contexts.jl

Co-authored-by: Hong Ge <[email protected]>
Signed-off-by: Will Tebbutt <[email protected]>

* Remove redundant line spacing

* Update src/interpreter/contexts.jl

Co-authored-by: Hong Ge <[email protected]>
Signed-off-by: Will Tebbutt <[email protected]>

* Explain the world age mechanism in is_primitive

* Improve docstring

* Fix some docstrings

* Docstring improved + comment added to implementation

* MinimalCtx is no longer a subtype of DefaultCtx.

* fix typos

* bugfix: fix dispatch error

* rm the `world` argument from `test_rule`

* Update Project.toml

Signed-off-by: Hong Ge <[email protected]>

---------

Signed-off-by: Will Tebbutt <[email protected]>
Signed-off-by: Hong Ge <[email protected]>
Co-authored-by: Hong Ge <[email protected]>
Co-authored-by: Hong Ge <[email protected]>

Author: willtebbutt

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.179"
+version = "0.4.180"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/developer_documentation/forwards_mode_design.md

@@ -70,8 +70,8 @@ Hand-written rules are implemented by writing methods of two functions: `is_prim
 
 ### `is_primitive`
 
-`is_primitive(::Type{<:Union{MinimalForwardsCtx, DefaultForwardsCtx}}, signature::Type{<:Tuple})` should return `true` if AD must attempt to differentiate a call by passing the arguments to `frule!!`, and `false` otherwise.
-The [`Mooncake.@is_primitive`](@ref) macro helps makes implementing this very easy.
+`is_primitive(::Type{<:Union{MinimalForwardsCtx, DefaultForwardsCtx}}, signature::Type{<:Tuple}, world)` must return `true` if AD must attempt to differentiate a call by passing the arguments to `frule!!`, and `false` otherwise.
+The [`Mooncake.@is_primitive`](@ref) macro must be used to extend to create new primitives.
 
 ### `frule!!`
 

docs/src/interface.md

@@ -7,8 +7,10 @@ See [Tutorial](@ref) for more info.
 
 ```@docs; canonical=true
 Mooncake.Config
+Mooncake.value_and_derivative!!
 Mooncake.value_and_gradient!!(::Mooncake.Cache, f::F, x::Vararg{Any, N}) where {F, N}
 Mooncake.value_and_pullback!!(::Mooncake.Cache, ȳ, f::F, x::Vararg{Any, N}) where {F, N}
+Mooncake.prepare_derivative_cache
 Mooncake.prepare_gradient_cache
 Mooncake.prepare_pullback_cache
 ```

src/Mooncake.jl

@@ -80,11 +80,12 @@ function rrule!! end
     build_primitive_rrule(sig::Type{<:Tuple})
 
 Construct an rrule for signature `sig`. For this function to be called in `build_rrule`, you
-must also ensure that `is_primitive(context_type, ReverseMode, sig)` is `true`. The callable
-returned by this must obey the rrule interface, but there are no restrictions on the type of
-callable itself. For example, you might return a callable `struct`. By default, this
-function returns `rrule!!` so, most of the time, you should just implement a method of
-`rrule!!`.
+must also ensure that a method of `_is_primitive(context_type, ReverseMode, sig)` exists,
+preferably by using the [@is_primitive](@ref) macro.
+The callable returned by this must obey the rrule interface, but there are no restrictions
+on the type of callable itself. For example, you might return a callable `struct`. By
+default, this function returns `rrule!!` so, most of the time, you should just implement a
+method of `rrule!!`.
 
 # Extended Help
 

src/interface.jl

@@ -594,7 +594,7 @@ function value_and_gradient!!(
 end
 
 """
-    prepare_derivative_cache(f, x...)
+    prepare_derivative_cache(fx...; kwargs...)
 
 Returns a cache used with [`value_and_derivative!!`](@ref). See that function for more info.
 """

src/interpreter/abstract_interpretation.jl

@@ -147,7 +147,11 @@ function Core.Compiler.abstract_call_gf_by_type(
         sv::CC.AbsIntState,
         max_methods::Int,
     )
-    is_primitive(C, M, atype) || return ret
+
+    # Check to see whether the call in question could possibly be a Mooncake primitive. If
+    # it could be, set its call info such that it will not be inlined away.
+    maybe_primitive(C, M, atype, interp.world) || return ret
+
     # Insert a `NoInlineCallInfo` to prevent any potential inlining.
     @static if VERSION < v"1.12-"
         call = ret::CC.CallMeta

src/interpreter/contexts.jl

@@ -40,31 +40,23 @@ function is a primitive in reverse-mode AD.
 struct ReverseMode <: Mode end
 
 """
-    is_primitive(::Type{Ctx}, ::Type{M}, sig) where {Ctx,M}
+    _is_primitive(context::Type, mode::Type{<:Mode}, sig::Type{<:Tuple})
 
-Returns a `Bool` specifying whether the methods specified by `sig` are considered primitives
-in the context of contexts of type `Ctx` in mode `M`.
-
-```julia
-is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(sin), Float64})
-```
-will return if calling `sin(5.0)` should be treated as primitive when the context is a
-`DefaultCtx`.
+This function is an internal implementation detail. It is used only by
+[`is_primitive`](@ref) and [`maybe_primitive`](@ref), and is used by these two functions in
+a very non-standard way. In particular, the value these functions return depends on the
+signatures of methods of this function, not what the methods do when invoked.
 
-Observe that this information means that whether or not something is a primitive in a
-particular context depends only on static information, not any run-time information that
-might live in a particular instance of `Ctx`.
+Generally speaking, you ought not to add methods to this function
+yourself, but make use of [`@is_primitive`](@ref).
 """
-is_primitive(::Type{MinimalCtx}, ::Type{<:Mode}, sig::Type{<:Tuple}) = false
-function is_primitive(::Type{DefaultCtx}, ::Type{M}, sig) where {M<:Mode}
-    return is_primitive(MinimalCtx, M, sig)
-end
+function _is_primitive end
 
 """
     @is_primitive context_type [mode_type] signature
 
-Creates a method of [`is_primitive`](@ref) which always returns `true` for the
-`context_type`, and `signature` provided. For example
+Declares that calls with signature `signature` are primitives in `context_type` and
+`mode_type`. For example
 ```jldoctest
 julia> using Mooncake: DefaultCtx, @is_primitive, is_primitive, ForwardMode, ReverseMode
 
@@ -73,16 +65,14 @@ foo (generic function with 1 method)
 
 julia> @is_primitive DefaultCtx Tuple{typeof(foo),Float64}
 
-julia> is_primitive(DefaultCtx, ForwardMode, Tuple{typeof(foo),Float64})
+julia> is_primitive(DefaultCtx, ForwardMode, Tuple{typeof(foo),Float64}, Base.get_world_counter())
 true
 
-julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Float64})
+julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Float64}, Base.get_world_counter())
 true
 ```
 Observe that this means that a rule is a primitive in all AD modes.
 
-You should implement more complicated methods of [`is_primitive`](@ref) in the usual way.
-
 Optionally, you can specify that a rule is only a primitive in a particular mode, eg.
 ```jldoctest
 julia> using Mooncake: DefaultCtx, @is_primitive, is_primitive, ForwardMode, ReverseMode
@@ -92,10 +82,10 @@ bar (generic function with 1 method)
 
 julia> @is_primitive DefaultCtx ForwardMode Tuple{typeof(bar),Float64}
 
-julia> is_primitive(DefaultCtx, ForwardMode, Tuple{typeof(bar),Float64})
+julia> is_primitive(DefaultCtx, ForwardMode, Tuple{typeof(bar),Float64}, Base.get_world_counter())
 true
 
-julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(bar),Float64})
+julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(bar),Float64}, Base.get_world_counter())
 false
 ```
 """
@@ -109,10 +99,173 @@ end
 
 function _is_primitive_expression(Tctx, Tmode, sig)
     return quote
-        function Mooncake.is_primitive(
+        function Mooncake._is_primitive(
             ::Type{$(esc(Tctx))}, ::Type{<:$(Tmode)}, ::Type{<:$(esc(sig))}
         )
             return true
         end
     end
 end
+
+const _IS_PRIMITIVE_CACHE_DefaultCtx = IdDict{Any,Bool}()
+const _IS_PRIMITIVE_CACHE_MinimalCtx = IdDict{Any,Bool}()
+
+"""
+    is_primitive(ctx::Type, mode::Type{<:Mode}, sig::Type{<:Tuple}, world::UInt)
+
+Returns a `Bool` specifying whether the methods specified by `sig` are considered primitives
+in the context of context `ctx` in mode `mode` at world age `world`.
+
+```jldoctest
+julia> using Mooncake: is_primitive, DefaultCtx, ReverseMode
+
+julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(sin), Float64}, Base.get_world_counter())
+true
+```
+
+`world` is needed as rules which Mooncake derives are associated to a particular Julia world
+age. As a result, anything declared a primitive after the construction of a rule ought not
+to be considered a primitive by that rule. One can explicitly derive a new rule (eg. via
+[`build_frule`](@ref), [`build_rrule`](@ref), or a function from the higher-level interface
+such as [`prepare_derivative_cache`](@ref), [`prepare_pullback_cache`](@ref) or
+[`prepare_gradient_cache`](@ref)) after new `@is_primitive` declarations, should it be
+needed in cases where a rule has been previously derived. To see how this works, consider
+the following:
+```jldoctest
+julia> using Mooncake: is_primitive, DefaultCtx, ReverseMode, @is_primitive
+
+julia> foo(x::Float64) = 5x
+foo (generic function with 1 method)
+
+julia> old_world_age = Base.get_world_counter();
+
+julia> @is_primitive DefaultCtx ReverseMode Tuple{typeof(foo),Float64}
+
+julia> new_world_age = Base.get_world_counter();
+
+julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Float64}, old_world_age)
+false
+
+julia> is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Float64}, new_world_age)
+true
+```
+Observe that `is_primitive` returns `false` for the world age prior to declaring `foo` a
+primitive, but `true` afterwards. For more information on Julia's world age mechanism, see
+https://docs.julialang.org/en/v1/manual/methods/#Redefining-Methods .
+"""
+function is_primitive(
+    ctx::Type{MinimalCtx}, mode::Type{<:Mode}, sig::Type{Tsig}, world::UInt
+) where {Tsig<:Tuple}
+    @nospecialize sig
+
+    # We don't ever need to evaluate this function for abstract `mode`s, and there is a
+    # performance penalty associated with doing so, so exclude the possibility.
+    isconcretetype(mode) || throw(ArgumentError("mode $mode is not a concrete type."))
+
+    # Check to see whether any methods of `_is_primitive` exist which apply to this
+    # ctx-mode-signature triple in world age `world`. If we have looked this up before,
+    # return the answer from the cache.
+
+    tt = Tuple{typeof(_is_primitive),Type{ctx},Type{mode},Type{sig}}
+    return get!(_IS_PRIMITIVE_CACHE_MinimalCtx, (world, tt)) do
+        return !isempty(Base._methods_by_ftype(tt, -1, world))
+    end
+end
+
+function is_primitive(
+    ctx::Type{DefaultCtx}, mode::Type{<:Mode}, sig::Type{Tsig}, world::UInt
+) where {Tsig<:Tuple}
+    @nospecialize sig
+
+    isconcretetype(mode) || throw(ArgumentError("mode $mode is not a concrete type."))
+
+    # This function returns `true` if the method is a primitive in either 
+    # `DefaultCtx` _or_ `MinimalCtx`. 
+    tt = Tuple{typeof(_is_primitive),Type{DefaultCtx},Type{mode},Type{sig}}
+    return get!(_IS_PRIMITIVE_CACHE_DefaultCtx, (world, tt)) do
+        return is_primitive(MinimalCtx, mode, sig, world) ||
+               !isempty(Base._methods_by_ftype(tt, -1, world))
+    end
+end
+
+const _MAYBE_PRIMITIVE_CACHE_MinimalCtx = IdDict{Any,Bool}()
+const _MAYBE_PRIMITIVE_CACHE_DefaultCtx = IdDict{Any,Bool}()
+
+"""
+    maybe_primitive(ctx::Type, mode::Type, sig::Type{<:Tuple}, world::UInt)
+
+`true` if there exists `M<:mode`, and `S<:sig` such that
+`is_primitive(ctx, M, S, world)` returns `true`.
+
+This functionality is used to determine whether or not it is safe to inline away a call
+site when performing abstract interpretation using a `MooncakeInterpreter`, which is only
+safe to do if the inferred argument types at the call site preclude the call being to a
+primitive.
+
+For example, consider the following:
+```jldoctest is_prim_example
+julia> using Mooncake: Mooncake, @is_primitive, DefaultCtx, ReverseMode
+
+julia> foo(x) = 5x;
+
+julia> @is_primitive DefaultCtx ReverseMode Tuple{typeof(foo),Float64}
+
+```
+This function agrees with [`is_primitive`](@ref) for fully inferred call sites:
+```jldoctest is_prim_example
+julia> world = Base.get_world_counter();
+
+julia> Mooncake.maybe_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Float64}, world)
+true
+
+julia> Mooncake.maybe_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Int}, world)
+false
+```
+However, it differs for call sites containing arguments whose types are not fully inferred.
+For example:
+```jldoctest is_prim_example
+julia> Mooncake.is_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Real}, world)
+false
+
+julia> Mooncake.maybe_primitive(DefaultCtx, ReverseMode, Tuple{typeof(foo),Real}, world)
+true
+```
+Per the definition at the top of this docstring, this function returns `true` because
+`Tuple{typeof(foo),Float64} <: Tuple{typeof(foo),Real}`.
+"""
+function maybe_primitive(
+    ctx::Type{MinimalCtx}, mode::Type{<:Mode}, sig::Type{Tsig}, world::UInt
+) where {Tsig<:Tuple}
+    @nospecialize sig
+
+    # We don't ever need to evaluate this function for abstract `mode`s, and there is a
+    # performance penalty associated with doing so, so exclude the possibility.
+    isconcretetype(mode) || throw(ArgumentError("mode $mode is not a concrete type."))
+
+    # Check to see whether any methods of `_is_primitive` exist which apply to any subtypes
+    # of this ctx-mode-signature triple in world age `world`. If we have looked this up
+    # before, return the answer from the cache.
+    tt = Tuple{typeof(_is_primitive),Type{ctx},Type{mode},Type{<:sig}}
+    return get!(_MAYBE_PRIMITIVE_CACHE_MinimalCtx, (world, tt)) do
+        return !isempty(Base._methods_by_ftype(tt, -1, world))
+    end
+end
+
+function maybe_primitive(
+    ctx::Type{DefaultCtx}, mode::Type{<:Mode}, sig::Type{Tsig}, world::UInt
+) where {Tsig<:Tuple}
+    @nospecialize sig
+
+    # We don't ever need to evaluate this function for abstract `mode`s, and there is a
+    # performance penalty associated with doing so, so exclude the possibility.
+    isconcretetype(mode) || throw(ArgumentError("mode $mode is not a concrete type."))
+
+    # Check to see whether any methods of `_is_primitive` exist which apply to any subtypes
+    # of this ctx-mode-signature triple in world age `world`. If we have looked this up
+    # before, return the answer from the cache.
+    tt = Tuple{typeof(_is_primitive),Type{ctx},Type{mode},Type{<:sig}}
+    return get!(_MAYBE_PRIMITIVE_CACHE_DefaultCtx, (world, tt)) do
+        return maybe_primitive(MinimalCtx, mode, sig, world) ||
+               !isempty(Base._methods_by_ftype(tt, -1, world))
+    end
+end

src/interpreter/forward_mode.jl

@@ -10,6 +10,17 @@ struct DualRuleInfo
     dual_ret_type::Type
 end
 
+"""
+    build_frule(
+        interp::MooncakeInterpreter{C},
+        sig_or_mi;
+        debug_mode=false,
+        silence_debug_messages=true,
+    ) where {C}
+
+Returns a function which performs forward-mode AD for `sig_or_mi`. Will derive a rule if
+`sig_or_mi` is not a primitive.
+"""
 function build_frule(
     interp::MooncakeInterpreter{C}, sig_or_mi; debug_mode=false, silence_debug_messages=true
 ) where {C}
@@ -33,7 +44,9 @@ function build_frule(
 
     # If we have a hand-coded rule, just use that.
     sig = _get_sig(sig_or_mi)
-    is_primitive(C, ForwardMode, sig) && return (debug_mode ? DebugFRule(frule!!) : frule!!)
+    if is_primitive(C, ForwardMode, sig, interp.world)
+        return (debug_mode ? DebugFRule(frule!!) : frule!!)
+    end
 
     # We don't have a hand-coded rule, so derive one.
     lock(MOONCAKE_INFERENCE_LOCK)
@@ -334,7 +347,8 @@ function modify_fwd_ad_stmts!(
             return uninit_dual(get_const_primal_value(arg))
         end
 
-        if is_primitive(context_type(info.interp), ForwardMode, sig)
+        interp = info.interp
+        if is_primitive(context_type(interp), ForwardMode, sig, interp.world)
             replace_call!(dual_ir, ssa, Expr(:call, frule!!, dual_args...))
         else
             dm = info.debug_mode
@@ -423,7 +437,7 @@ function frule_type(
     interp::MooncakeInterpreter{C}, mi::CC.MethodInstance; debug_mode
 ) where {C}
     primal_sig = _get_sig(mi)
-    if is_primitive(C, ForwardMode, primal_sig)
+    if is_primitive(C, ForwardMode, primal_sig, interp.world)
         return debug_mode ? DebugFRule{typeof(frule!!)} : typeof(frule!!)
     end
     ir, _ = lookup_ir(interp, mi)

src/interpreter/reverse_mode.jl

@@ -702,7 +702,8 @@ function make_ad_stmts!(stmt::Expr, line::ID, info::ADInfo)
 
         # Construct signature, and determine how the rrule is to be computed.
         sig = Tuple{arg_types...}
-        raw_rule = if is_primitive(context_type(info.interp), ReverseMode, sig)
+        interp = info.interp
+        raw_rule = if is_primitive(context_type(interp), ReverseMode, sig, interp.world)
             rrule!! # intrinsic / builtin / thing we provably have rule for
         elseif is_invoke
             mi = get_mi(stmt.args[1])
@@ -1123,7 +1124,7 @@ function build_rrule(
 
     # If we have a hand-coded rule, just use that.
     sig = _get_sig(sig_or_mi)
-    if is_primitive(C, ReverseMode, sig)
+    if is_primitive(C, ReverseMode, sig, interp.world)
         rule = build_primitive_rrule(sig)
         return (debug_mode ? DebugRRule(rule) : rule)
     end
@@ -1900,7 +1901,7 @@ important for performance in dynamic dispatch, and to ensure that recursion work
 properly.
 """
 function rule_type(interp::MooncakeInterpreter{C}, sig_or_mi; debug_mode) where {C}
-    if is_primitive(C, ReverseMode, _get_sig(sig_or_mi))
+    if is_primitive(C, ReverseMode, _get_sig(sig_or_mi), interp.world)
         rule = build_primitive_rrule(_get_sig(sig_or_mi))
         return debug_mode ? DebugRRule{typeof(rule)} : typeof(rule)
     end