Start forward mode AD

src/Mooncake.jl

@@ -22,6 +22,7 @@ using Base:
     arrayset, TwicePrecision, twiceprecision
 using Base.Experimental: @opaque
 using Base.Iterators: product
+using Base.Meta: isexpr
 using Core:
     Intrinsics, bitcast, SimpleVector, svec, ReturnNode, GotoNode, GotoIfNot, PhiNode,
     PiNode, SSAValue, Argument, OpaqueClosure, compilerbarrier
@@ -34,6 +35,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...)
 
@@ -61,6 +69,7 @@ function rrule!! end
 
 include("utils.jl")
 include("tangents.jl")
+include("dual.jl")
 include("fwds_rvs_data.jl")
 include("codual.jl")
 include("debug_mode.jl")
@@ -72,12 +81,15 @@ include(joinpath("interpreter", "ir_utils.jl"))
 include(joinpath("interpreter", "bbcode.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")
 include("test_utils.jl")
 include("test_resources.jl")
 
+include(joinpath("frules", "basic.jl"))
+
 include(joinpath("rrules", "avoiding_non_differentiable_code.jl"))
 include(joinpath("rrules", "blas.jl"))
 include(joinpath("rrules", "builtins.jl"))
@@ -118,9 +130,13 @@ export
     _add_to_primal,
     _diff,
     _dot,
+    Dual,
+    zero_dual,
     zero_codual,
     codual_type,
+    frule!!,
     rrule!!,
+    build_frule,
     build_rrule,
     value_and_gradient!!,
     value_and_pullback!!,

src/debug_mode.jl

@@ -1,3 +1,4 @@
+DebugFRule(rule) = rule  # TODO: make it non-trivial
 
 """
     DebugPullback(pb, y, x)

src/dual.jl

@@ -0,0 +1,13 @@
+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
+
+zero_dual(x) = Dual(x, zero_tangent(x))
+
+dual_type(::Type{P}) where {P} = Dual{P,tangent_type(P)}

src/frules/basic.jl

@@ -0,0 +1,11 @@
+frule!!(f::F, args::Vararg{Dual,N}) where {F,N} = frule!!(zero_dual(f), args...)
+
+@is_primitive MinimalCtx Tuple{typeof(sin),Number}
+function frule!!(::Dual{typeof(sin)}, x::Dual{<:Number})
+    return Dual(sin(primal(x)), cos(primal(x)) * tangent(x))
+end
+
+@is_primitive MinimalCtx Tuple{typeof(cos),Number}
+function frule!!(::Dual{typeof(cos)}, x::Dual{<:Number})
+    return Dual(cos(primal(x)), -sin(primal(x)) * tangent(x))
+end

src/interpreter/s2s_forward_mode_ad.jl

@@ -0,0 +1,147 @@
+function build_frule(args...; debug_mode=false)
+    interp = get_interpreter()
+    sig = _typeof(TestUtils.__get_primals(args))
+    return build_frule(interp, sig; debug_mode)
+end
+
+function build_frule(
+    interp::MooncakeInterpreter{C},
+    sig_or_mi;
+    debug_mode=false,
+    silence_debug_messages=true,
+) where {C}
+    # To avoid segfaults, ensure that we bail out if the interpreter's world age is greater
+    # than the current world age.
+    if Base.get_world_counter() > interp.world
+        throw(
+            ArgumentError(
+                "World age associated to interp is behind current world age. Please " *
+                "a new interpreter for the current world age.",
+            ),
+        )
+    end
+
+    # If we're compiling in debug mode, let the user know by default.
+    if !silence_debug_messages && debug_mode
+        @info "Compiling rule for $sig_or_mi in debug mode. Disable for best performance."
+    end
+
+    # If we have a hand-coded rule, just use that.
+    _is_primitive(C, sig_or_mi) && return (debug_mode ? DebugFRule(frule!!) : frule!!)
+
+
+    # We don't have a hand-coded rule, so derived one.
+    lock(MOONCAKE_INFERENCE_LOCK)
+    try
+        # If we've already derived the OpaqueClosures and info, do not re-derive, just
+        # create a copy and pass in new shared data.
+        oc_cache_key = ClosureCacheKey(interp.world, (sig_or_mi, debug_mode))
+        if haskey(interp.oc_cache, oc_cache_key)
+            return _copy(interp.oc_cache[oc_cache_key])
+        else
+            # Derive forward-pass IR, and shove in a `MistyClosure`.
+            forward_ir = generate_forward_ir(interp, sig_or_mi; debug_mode)
+            fwd_oc = MistyClosure(forward_ir; do_compile=true)
+            raw_rule = DerivedFRule(fwd_oc)
+            rule = debug_mode ? DebugFRule(raw_rule) : raw_rule
+            interp.oc_cache[oc_cache_key] = rule
+            return rule
+        end
+    catch e
+        rethrow(e)
+    finally
+        unlock(MOONCAKE_INFERENCE_LOCK)
+    end
+end
+
+function generate_forward_ir(
+    interp::MooncakeInterpreter,
+    sig_or_mi;
+    debug_mode=false,
+    do_inline=true,
+)
+    # Reset id count. This ensures that the IDs generated are the same each time this
+    # function runs.
+    seed_id!()
+
+    # Grab code associated to the primal.
+    primal_ir, _ = lookup_ir(interp, sig_or_mi)
+
+    # Normalise the IR.
+    isva, spnames = is_vararg_and_sparam_names(sig_or_mi)
+    ir = normalise!(primal_ir, spnames)
+
+    fwd_ir = dualize_ir(ir)
+    opt_fwd_ir = optimise_ir!(fwd_ir; do_inline)
+    return opt_fwd_ir
+end
+
+function dualize_ir(ir::IRCode)
+    new_stmts_stmt = map(make_fwd_ad_stmt, ir.stmts.stmt)
+    new_stmts_type = map(dual_type, ir.stmts.type)
+    new_stmts_info = ir.stmts.info
+    new_stmts_line = ir.stmts.line
+    new_stmts_flag = ir.stmts.flag
+    new_stmts = CC.InstructionStream(
+        new_stmts_stmt,
+        new_stmts_type,
+        new_stmts_info,
+        new_stmts_line,
+        new_stmts_flag,
+    )
+    new_cfg = ir.cfg
+    new_linetable = ir.linetable
+    rule_type = Any
+    new_argtypes = convert(Vector{Any}, vcat(rule_type, map(make_fwd_argtype, ir.argtypes)))
+    new_meta = ir.meta
+    new_sptypes = ir.sptypes
+    return IRCode(new_stmts, new_cfg, new_linetable, new_argtypes, new_meta, new_sptypes)
+end
+
+make_fwd_argtype(::Type{P}) where {P} = dual_type(P)
+make_fwd_argtype(c::Core.Const) = Dual  # TODO: refine to type of const
+
+function make_fwd_ad_stmt(stmt::Expr)
+    interp = get_interpreter()  # TODO: pass it around
+    C = context_type(interp)
+    if isexpr(stmt, :invoke) || isexpr(stmt, :call)
+        mi = stmt.args[1]::Core.MethodInstance
+        sig = mi.specTypes
+        if is_primitive(C, sig)
+            shifted_args = map(stmt.args) do a
+                if a isa Core.Argument
+                    Core.Argument(a.n + 1)
+                else
+                    a
+                end
+            end
+            new_stmt = Expr(
+                :call,
+                :($frule!!),
+                stmt.args[2],
+                shifted_args[3:end]...
+            )
+            return new_stmt
+        else
+            throw(ArgumentError("Recursing into non-primitive calls is not yet supported in forward mode"))
+        end
+        return stmt
+    else
+        throw(ArgumentError("Expressions of type `:$(stmt.head)` are not yet supported in forward mode"))
+    end
+    return stmt
+end
+
+function make_fwd_ad_stmt(stmt::ReturnNode)
+    return stmt
+end
+
+struct DerivedFRule{Tfwd_oc}
+    fwd_oc::Tfwd_oc
+end
+
+_copy(rule::DerivedFRule) = deepcopy(rule)
+
+@inline function (fwd::DerivedFRule)(args::Vararg{Dual,N}) where {N}
+    return fwd.fwd_oc.oc(args...)
+end

src/test_utils.jl

@@ -470,7 +470,7 @@ function test_rrule_performance(
     end
 end
 
-__get_primals(xs) = map(x -> x isa CoDual ? primal(x) : x, xs)
+__get_primals(xs) = map(x -> x isa Union{Dual,CoDual} ? primal(x) : x, xs)
 
 @doc"""
     test_rule(

test/forward.jl

@@ -0,0 +1,27 @@
+using Mooncake
+using Test
+
+x, dx = 2.0, 3.0
+xdual = Dual(x, dx)
+
+@testset "Manual frule" begin
+    sin_rule = build_frule(sin, x)
+    ydual = sin_rule(zero_dual(sin), xdual)
+
+    @test primal(ydual) == sin(x)
+    @test tangent(ydual) == dx * cos(x)
+end
+
+function func(x)
+    y = sin(x)
+    z = cos(y)
+    return z
+end
+
+@testset "Automatic frule" begin
+    func_rule = build_frule(func, x)
+    ydual = func_rule(zero_dual(func), xdual)
+
+    @test primal(ydual) == cos(sin(x))
+    @test tangent(ydual) ≈ dx * -sin(sin(x)) * cos(x)
+end

test/runtests.jl

@@ -24,6 +24,7 @@ include("front_matter.jl")
         include("config.jl")
         include("developer_tools.jl")
         include("test_utils.jl")
+        include("forward.jl")
     elseif test_group == "rrules/avoiding_non_differentiable_code"
         include(joinpath("rrules", "avoiding_non_differentiable_code.jl"))
     elseif test_group == "rrules/blas"