Forward-mode AD

.github/workflows/CI.yml

@@ -41,6 +41,7 @@ jobs:
           'rrules/low_level_maths',
           'rrules/memory',
           'rrules/misc',
+          'rrules/misty_closures',
           'rrules/new',
           'rrules/random',
           'rrules/tasks',

HISTORY.md

@@ -0,0 +1,11 @@
+# 0.4.143
+
+## Public Interface
+- Mooncake offers forward mode AD.
+- Two new functions added to the public interface: `prepare_derivative_cache` and `value_and_derivative!!`.
+- One new type added to the public interface: `Dual`.
+
+## Internals
+- `get_interpreter` was previously a zero-arg function. Is now a unary function, called with a "mode" argument: `get_interpreter(ForwardMode)`, `get_interpreter(ReverseMode)`.
+- `@zero_derivative` should now be preferred to `@zero_adjoint`. `@zero_adjoint` will be removed in 0.5.
+- `@from_chainrules` should now be preferred to `@from_rrule`. `@from_rrule` will be removed in 0.5.

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.146"
+version = "0.4.147"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -59,8 +59,9 @@ FunctionWrappers = "1.1.3"
 GPUArraysCore = "0.1, 0.2"
 Graphs = "1"
 InteractiveUtils = "1"
-JET = "0.9, 0.10"
+JET = "0.9"
 JuliaFormatter = "1.0, 2.1"
+JuliaInterpreter = "0.9"
 LinearAlgebra = "1"
 LuxLib = "1"
 MistyClosures = "2"
@@ -81,9 +82,10 @@ DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "JuliaFormatter", "Pkg", "StableRNGs", "Test"]
+test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "JuliaFormatter", "JuliaInterpreter", "Pkg", "StableRNGs", "Test"]

bench/run_benchmarks.jl

@@ -19,17 +19,21 @@ using AbstractGPs,
     Zygote
 
 using Mooncake:
+    Dual,
     CoDual,
     generate_hand_written_rrule!!_test_cases,
     generate_derived_rrule!!_test_cases,
     TestUtils,
     _typeof,
     primal,
     tangent,
+    zero_dual,
     zero_codual
 
 using Mooncake.TestUtils: _deepcopy
 
+to_benchmark(__frule!!::R, dx::Vararg{Dual,N}) where {R,N} = __frule!!(dx...)
+
 function to_benchmark(__rrule!!::R, dx::Vararg{CoDual,N}) where {R,N}
     dx_f = Mooncake.tuple_map(x -> CoDual(primal(x), Mooncake.fdata(tangent(x))), dx)
     out, pb!! = __rrule!!(dx_f...)
@@ -206,6 +210,20 @@ function benchmark_rules!!(test_case_data, default_ratios, include_other_framewo
                 evals=1,
             )
 
+            # Benchmark AD via Mooncake.
+            @info "Mooncake (Forward)"
+            rule = Mooncake.build_frule(args...)
+            duals = map(x -> x isa CoDual ? Dual(x.x, x.dx) : zero_dual(x), args)
+            to_benchmark(rule, duals...)
+            include_other_frameworks && GC.gc(true)
+            suite["mooncake_fwd"] = Chairmarks.benchmark(
+                () -> (rule, duals),
+                identity,
+                a -> to_benchmark(a[1], a[2]...),
+                _ -> true;
+                evals=1,
+            )
+
             if include_other_frameworks
                 if should_run_benchmark(Val(:zygote), args...)
                     @info "Zygote"
@@ -258,6 +276,7 @@ function combine_results(result, tag, _range, default_range)
     d = result[2]
     primal_time = minimum(d["primal"]).time
     mooncake_time = minimum(d["mooncake"]).time
+    mooncake_fwd_time = minimum(d["mooncake_fwd"]).time
     zygote_time = in("zygote", keys(d)) ? minimum(d["zygote"]).time : missing
     rd_time = in("rd", keys(d)) ? minimum(d["rd"]).time : missing
     ez_time = in("enzyme", keys(d)) ? minimum(d["enzyme"]).time : missing
@@ -267,6 +286,8 @@ function combine_results(result, tag, _range, default_range)
         primal_time=primal_time,
         mooncake_time=mooncake_time,
         Mooncake=mooncake_time / primal_time,
+        mooncake_fwd_time=mooncake_fwd_time,
+        MooncakeFwd=mooncake_fwd_time / primal_time,
         zygote_time=zygote_time,
         Zygote=zygote_time / primal_time,
         rd_time=rd_time,
@@ -348,7 +369,7 @@ end
 
 function create_inter_ad_benchmarks()
     results = benchmark_inter_framework_rules()
-    tools = [:Mooncake, :Zygote, :ReverseDiff, :Enzyme]
+    tools = [:Mooncake, :MooncakeFwd, :Zygote, :ReverseDiff, :Enzyme]
     df = DataFrame(results)[:, [:tag, :primal_time, tools...]]
 
     # Plot graph of results.

ext/MooncakeSpecialFunctionsExt.jl

@@ -3,44 +3,44 @@ module MooncakeSpecialFunctionsExt
 using SpecialFunctions, Mooncake
 using Base: IEEEFloat
 
-import Mooncake: @from_rrule, DefaultCtx, @zero_adjoint
+import Mooncake: DefaultCtx, @from_chainrules, @zero_derivative
 
-@from_rrule DefaultCtx Tuple{typeof(airyai),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(airyaix),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(airyaiprime),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(airybi),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(airybiprime),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(besselj0),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(besselj1),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(bessely0),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(bessely1),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(dawson),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(digamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erf),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erf),IEEEFloat,IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erfc),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(logerfc),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erfcinv),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erfcx),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(logerfcx),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erfi),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(erfinv),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(gamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(invdigamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(trigamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(polygamma),Integer,IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(beta),IEEEFloat,IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(logbeta),IEEEFloat,IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(logabsgamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(loggamma),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(expint),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(expintx),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(expinti),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(sinint),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(cosint),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(ellipk),IEEEFloat}
-@from_rrule DefaultCtx Tuple{typeof(ellipe),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(airyai),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(airyaix),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(airyaiprime),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(airybi),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(airybiprime),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(besselj0),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(besselj1),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(bessely0),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(bessely1),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(dawson),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(digamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erf),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erf),IEEEFloat,IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erfc),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(logerfc),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erfcinv),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erfcx),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(logerfcx),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erfi),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(erfinv),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(gamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(invdigamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(trigamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(polygamma),Integer,IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(beta),IEEEFloat,IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(logbeta),IEEEFloat,IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(logabsgamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(loggamma),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(expint),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(expintx),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(expinti),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(sinint),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(cosint),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(ellipk),IEEEFloat}
+@from_chainrules DefaultCtx Tuple{typeof(ellipe),IEEEFloat}
 
-@zero_adjoint DefaultCtx Tuple{typeof(logfactorial),Integer}
+@zero_derivative DefaultCtx Tuple{typeof(logfactorial),Integer}
 
 end

src/Mooncake.jl

@@ -19,6 +19,7 @@ using Base:
     twiceprecision
 using Base.Experimental: @opaque
 using Base.Iterators: product
+using Base.Meta: isexpr
 using Core:
     Intrinsics,
     bitcast,
@@ -29,6 +30,8 @@ using Core:
     GotoIfNot,
     PhiNode,
     PiNode,
+    PhiCNode,
+    UpsilonNode,
     SSAValue,
     Argument,
     OpaqueClosure,
@@ -43,6 +46,13 @@ using DispatchDoctor: @stable, @unstable
 # Needs to be defined before various other things.
 function _foreigncall_ end
 
+"""
+    frule!!(f::Dual, x::Dual...)
+
+Performs AD in forward mode, possibly modifying the inputs, and returns a `Dual`.
+"""
+function frule!! end
+
 """
     rrule!!(f::CoDual, x::CoDual...)
 
@@ -72,10 +82,11 @@ 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, 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 `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!!`.
 
 # Extended Help
 
@@ -95,6 +106,7 @@ build_primitive_rrule(::Type{<:Tuple}) = rrule!!
 @stable default_mode = "disable" default_union_limit = 2 begin
 include("utils.jl")
 include("tangents.jl")
+include("dual.jl")
 include("fwds_rvs_data.jl")
 include("codual.jl")
 include("debug_mode.jl")
@@ -110,7 +122,8 @@ include(joinpath("interpreter", "patch_for_319.jl"))
 include(joinpath("interpreter", "ir_utils.jl"))
 include(joinpath("interpreter", "ir_normalisation.jl"))
 include(joinpath("interpreter", "zero_like_rdata.jl"))
-include(joinpath("interpreter", "s2s_reverse_mode_ad.jl"))
+include(joinpath("interpreter", "forward_mode.jl"))
+include(joinpath("interpreter", "reverse_mode.jl"))
 end
 
 include("tools_for_rules.jl")
@@ -129,6 +142,7 @@ include(joinpath("rrules", "lapack.jl"))
 include(joinpath("rrules", "linear_algebra.jl"))
 include(joinpath("rrules", "low_level_maths.jl"))
 include(joinpath("rrules", "misc.jl"))
+include(joinpath("rrules", "misty_closures.jl"))
 include(joinpath("rrules", "new.jl"))
 include(joinpath("rrules", "random.jl"))
 include(joinpath("rrules", "tasks.jl"))
@@ -146,12 +160,15 @@ include("developer_tools.jl")
 
 # Public, not exported
 include("public.jl")
+
 end
 #! format: on
 
-@public Config, value_and_pullback!!, prepare_pullback_cache
+@public Config, value_and_pullback!!, prepare_pullback_cache, value_and_derivative!!
+@public prepare_derivative_cache, Dual
 
 # Public, exported
-export value_and_gradient!!, prepare_gradient_cache
+export value_and_gradient!!, prepare_gradient_cache, value_and_derivative!!
+export prepare_derivative_cache
 
 end

src/codual.jl

@@ -17,6 +17,13 @@ tangent(x::CoDual) = x.dx
 Base.copy(x::CoDual) = CoDual(copy(primal(x)), copy(tangent(x)))
 _copy(x::P) where {P<:CoDual} = x
 
+"""
+    extract(x::CoDual)
+
+Helper function. Returns the 2-tuple `x.x, x.dx`.
+"""
+extract(x::CoDual) = primal(x), tangent(x)
+
 """
     zero_codual(x)
 

src/debug_mode.jl

@@ -1,3 +1,5 @@
+# TODO: make it non-trivial. See https://github.com/chalk-lab/Mooncake.jl/issues/672
+DebugFRule(rule) = rule
 
 """
     DebugPullback(pb, y, x)