Start forward mode AD

.buildkite/pipeline.yml

@@ -23,7 +23,7 @@ steps:
       setup:
         version:
           - "1"
-          - "1.10"
+          # - "1.10"
         label:
           - "cuda"
           - "nnlib"

.github/workflows/CI.yml

@@ -33,7 +33,7 @@ jobs:
           'rrules/builtins',
           'rrules/fastmath',
           'rrules/foreigncall',
-          'rrules/functionwrappers',
+          # 'rrules/functionwrappers',
           'rrules/iddict',
           'rrules/lapack',
           'rrules/linear_algebra',

.gitignore

@@ -6,3 +6,4 @@ profile.pb.gz
 scratch.jl
 docs/build/
 docs/site/
+playground.jl

Project.toml

@@ -54,8 +54,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"
@@ -75,9 +76,10 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 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...)
@@ -207,6 +211,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 Dual ? x : 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"
@@ -259,6 +277,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
@@ -268,6 +287,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,
@@ -349,7 +370,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_chain_rule, @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_chain_rule DefaultCtx Tuple{typeof(airyai),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(airyaix),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(airyaiprime),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(airybi),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(airybiprime),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(besselj0),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(besselj1),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(bessely0),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(bessely1),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(dawson),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(digamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erf),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erf),IEEEFloat,IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erfc),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(logerfc),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erfcinv),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erfcx),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(logerfcx),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erfi),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(erfinv),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(gamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(invdigamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(trigamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(polygamma),Integer,IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(beta),IEEEFloat,IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(logbeta),IEEEFloat,IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(logabsgamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(loggamma),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(expint),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(expintx),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(expinti),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(sinint),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(cosint),IEEEFloat}
+@from_chain_rule DefaultCtx Tuple{typeof(ellipk),IEEEFloat}
+@from_chain_rule 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,
@@ -42,6 +43,13 @@ using FunctionWrappers: FunctionWrapper
 # 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...)
 
@@ -92,6 +100,7 @@ build_primitive_rrule(::Type{<:Tuple}) = rrule!!
 
 include("utils.jl")
 include("tangents.jl")
+include("dual.jl")
 include("fwds_rvs_data.jl")
 include("codual.jl")
 include("debug_mode.jl")
@@ -106,6 +115,7 @@ 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_forward_mode_ad.jl"))
 include(joinpath("interpreter", "s2s_reverse_mode_ad.jl"))
 
 include("tools_for_rules.jl")
@@ -123,6 +133,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"))

src/codual.jl

@@ -17,6 +17,13 @@
 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,4 @@
+DebugFRule(rule) = rule  # TODO: make it non-trivial
 
 """
     DebugPullback(pb, y, x)

src/developer_tools.jl

@@ -104,3 +104,15 @@
 )::IRCode
     return generate_ir(interp, sig; debug_mode, do_inline).rvs_ir
 end
+
+"""
+
+"""
+function dual_ir(
+    sig::Type{<:Tuple};
+    interp=get_interpreter(),
+    debug_mode::Bool=false,
+    do_inline::Bool=true,
+)
+    return generate_dual_ir(interp, sig; debug_mode, do_inline)
+end

src/dual.jl

@@ -0,0 +1,47 @@
+struct Dual{P,T}
+    primal::P
+    tangent::T
+end
+
+primal(x::Dual) = x.primal
+tangent(x::Dual) = x.tangent
+Base.copy(x::Dual) = Dual(copy(primal(x)), copy(tangent(x)))
+_copy(x::P) where {P<:Dual} = x
+
+"""
+    extract(x::CoDual)
+
+Helper function. Returns the 2-tuple `x.x, x.dx`.
+"""
+extract(x::Dual) = primal(x), tangent(x)
+
+zero_dual(x) = Dual(x, zero_tangent(x))
+randn_dual(rng::AbstractRNG, x) = Dual(x, randn_tangent(rng, x))
+
+function dual_type(::Type{P}) where {P}
+    P == DataType && return Dual
+    P isa Union && return Union{dual_type(P.a),dual_type(P.b)}
+    P <: UnionAll && return Dual # P is abstract, so we don't know its tangent type.
+    return isconcretetype(P) ? Dual{P,tangent_type(P)} : Dual
+end
+
+function dual_type(p::Type{Type{P}}) where {P}
+    return @isdefined(P) ? Dual{Type{P},NoTangent} : Dual{_typeof(p),NoTangent}
+end
+
+_primal(x) = x
+_primal(x::Dual) = primal(x)
+
+"""
+    verify_dual_type(x::Dual)
+
+Check that the type of `tangent(x)` is the tangent type of the type of `primal(x)`.
+"""
+verify_dual_type(x::Dual) = tangent_type(typeof(primal(x))) == typeof(tangent(x))
+
+@inline uninit_dual(x::P) where {P} = Dual(x, uninit_tangent(x))
+
+# Always sharpen the first thing if it's a type so static dispatch remains possible.
+function Dual(x::Type{P}, dx::NoTangent) where {P}
+    return Dual{@isdefined(P) ? Type{P} : typeof(x),NoTangent}(x, dx)
+end

src/interpreter/ir_normalisation.jl

@@ -63,6 +63,7 @@
         if stmt isa Expr && stmt.head == :boundscheck && length(stmt.args) == 1
             def = SSAValue(n)
             val = only(stmt.args)
+            # TODO: this could just be `statements[n] = val` (Valentin C says)
             for (m, stmt) in enumerate(statements)
                 statements[m] = replace_uses_with!(stmt, def, val)
             end
@@ -349,6 +350,7 @@
 
 @is_primitive MinimalCtx Tuple{typeof(gc_preserve),Vararg{Any,N}} where {N}
 
+frule!!(::Dual{typeof(gc_preserve)}, ::Vararg{Dual,N}) where {N} = zero_dual(nothing)
 function rrule!!(f::CoDual{typeof(gc_preserve)}, xs::CoDual...)
     pb = NoPullback(f, xs...)
     gc_preserve_pb!!(::NoRData) = GC.@preserve xs pb(NoRData())