Try to force better forward mode inference (#111)

By adding explicit recursion functions, and reducing the amount of
inference work required for higher order forward mode.

Author: Keno

src/Diffractor.jl

@@ -13,7 +13,6 @@ include("tangent.jl")
 include("jet.jl")
 
 include("stage1/generated.jl")
-include("stage1/termination.jl")
 include("stage1/forward.jl")
 include("stage1/recurse_fwd.jl")
 include("stage1/mixed.jl")
@@ -36,4 +35,6 @@ include("higher_fwd_rules.jl")
 
 include("debugutils.jl")
 
+include("stage1/termination.jl")
+
 end

src/stage1/forward.jl

@@ -10,36 +10,27 @@ primal(z::ZeroTangent) = ZeroTangent()
 
 first_partial(x) = partial(x, 1)
 
-# TODO: Which version do we want in ChainRules?
-function my_frule(args::ATB{1}...)
-    frule(DiffractorRuleConfig(), map(first_partial, args), map(primal, args)...)
-end
-
-# Fast path for some hot cases
-my_frule(::ZeroBundle{1, typeof(frule)}, args::ATB{1}...) = nothing
-my_frule(::ZeroBundle{1, typeof(my_frule)}, args::ATB{1}...) = nothing
-
-(::∂☆{N})(::ZeroBundle{N, typeof(my_frule)}, ::ZeroBundle{N, ZeroBundle{1, typeof(frule)}}, args::ATB{N}...) where {N} = ZeroBundle{N}(nothing)
-(::∂☆{N})(::ZeroBundle{N, typeof(my_frule)}, ::ZeroBundle{N, ZeroBundle{1, typeof(my_frule)}}, args::ATB{N}...) where {N} = ZeroBundle{N}(nothing)
-
 shuffle_down(b::UniformBundle{N, B, U}) where {N, B, U} =
-    UniformBundle{minus1(N), <:Any, U}(UniformBundle{1, B, U}(b.primal, b.tangent.val), b.tangent.val)
+    UniformBundle{minus1(N), <:Any}(UniformBundle{1, B}(b.primal, b.tangent.val),
+                                    UniformBundle{1, U}(b.tangent.val, b.tangent.val))
 
 function shuffle_down(b::ExplicitTangentBundle{N, B}) where {N, B}
     # N.B: This depends on the special properties of the canonical tangent index order
+    Base.@constprop :aggressive function _sdown(i::Int64)
+        ExplicitTangentBundle{1}(partial(b, 2*i), (partial(b, 2*i+1),))
+    end
     ExplicitTangentBundle{N-1}(
         ExplicitTangentBundle{1}(b.primal, (partial(b, 1),)),
-        ntuple(2^(N-1)-1) do i
-            ExplicitTangentBundle{1}(partial(b, 2*i), (partial(b, 2*i+1),))
-        end)
+        ntuple(_sdown, 2^(N-1)-1))
 end
 
 function shuffle_down(b::TaylorBundle{N, B}) where {N, B}
+    Base.@constprop :aggressive function _sdown(i::Int64)
+        ExplicitTangentBundle{1}(b.tangent.coeffs[i], (b.tangent.coeffs[i+1],))
+    end
     TaylorBundle{N-1}(
         ExplicitTangentBundle{1}(b.primal, (b.tangent.coeffs[1],)),
-        ntuple(N-1) do i
-            ExplicitTangentBundle{1}(b.tangent.coeffs[i], (b.tangent.coeffs[i+1],))
-        end)
+        ntuple(_sdown, N-1))
 end
 
 function shuffle_down(b::CompositeBundle{N, B}) where {N, B}
@@ -106,10 +97,17 @@ end
 struct ∂☆internal{N}; end
 struct ∂☆shuffle{N}; end
 
-shuffle_base(r) = TaylorBundle{1}(r[1], (r[2],))
+function shuffle_base(r)
+    (primal, dual) = r
+    if isa(dual, Union{NoTangent, ZeroTangent})
+        UniformBundle{1}(primal, dual)
+    else
+        TaylorBundle{1}(primal, (dual,))
+    end
+end
 
 function (::∂☆internal{1})(args::AbstractTangentBundle{1}...)
-    r = my_frule(args...)
+    r = frule(#=DiffractorRuleConfig(),=# map(first_partial, args), map(primal, args)...)
     if r === nothing
         return ∂☆recurse{1}()(args...)
     else
@@ -125,7 +123,9 @@ end
 
 function (::∂☆shuffle{N})(args::AbstractTangentBundle{N}...) where {N}
     ∂☆p = ∂☆{minus1(N)}()
-    ∂☆p(ZeroBundle{minus1(N)}(my_frule), map(shuffle_down, args)...)
+    downargs = map(shuffle_down, args)
+    tupargs = ∂vararg{minus1(N)}()(map(first_partial, downargs)...)
+    ∂☆p(ZeroBundle{minus1(N)}(frule), #= ZeroBundle{minus1(N)}(DiffractorRuleConfig()), =# tupargs, map(primal, downargs)...)
 end
 
 function (::∂☆internal{N})(args::AbstractTangentBundle{N}...) where {N}

src/stage1/recurse.jl

@@ -261,7 +261,7 @@ function transform!(ci, meth, nargs, sparams, N)
     cfg = compute_basic_blocks(code)
     slotnames = Symbol[Symbol("#self#"), :args, ci.slotnames...]
     slotflags = UInt8[(0x00 for i = 1:2)..., ci.slotflags...]
-    slottypes = UInt8[(0x00 for i = 1:2)..., ci.slotflags...]
+    slottypes = ci.slottypes === nothing ? nothing : UInt8[(Any for i = 1:2)..., ci.slottypes...]
 
     meta = Expr[]
     ir = IRCode(Core.Compiler.InstructionStream(code, Any[],

src/stage1/recurse_fwd.jl

@@ -54,7 +54,7 @@ function perform_fwd_transform(@nospecialize(ff::Type{∂☆recurse{N}}), @nospe
     ci′.method_for_inference_limit_heuristics = match.method
     slotnames = Symbol[Symbol("#self#"), :args, ci.slotnames...]
     slotflags = UInt8[(0x00 for i = 1:2)..., ci.slotflags...]
-    slottypes = Any[(Any for i = 1:2)..., ci.slotflags...]
+    slottypes = ci.slottypes === nothing ? nothing : Any[(Any for i = 1:2)..., ci.slottypes...]
     ci′.slotnames = slotnames
     ci′.slotflags = slotflags
     ci′.slottypes = slottypes

src/stage1/termination.jl

@@ -44,4 +44,30 @@ which(Tuple{∂⃖{N}, ∂⃖{1}, Vararg{Any}} where {N}).recursion_relation = f
 	       isa(Base.unwrap_unionall(new_sig.parameters[1].parameters[1]), Int)
 end
 
+for (;method) in Base._methods_by_ftype(Tuple{Diffractor.∂☆recurse{N}, Vararg{Any}} where {N}, nothing, -1, typemax(UInt64))
+	method.recursion_relation = function (method1, method2, parent_sig, new_sig)
+		# Recursion from a higher to a lower order is always allowed
+		parent_order = parent_sig.parameters[1].parameters[1]
+		child_order = new_sig.parameters[1].parameters[1]
+		#@Core.Main.Base.show (parent_order, child_order)
+		if parent_order > child_order
+			return true
+		end
+		@show (parent_sig, new_sig)
+		return false
+	end
+end
+
+for (;method) in Base._methods_by_ftype(Tuple{Diffractor.∂☆internal{N}, Vararg{Any}} where {N}, nothing, -1, typemax(UInt64))
+	method.recursion_relation = function (method1, method2, parent_sig, new_sig)
+		return true
+	end
+end
+
+for (;method) in Base._methods_by_ftype(Tuple{Diffractor.∂☆{N}, Vararg{Any}} where {N}, nothing, -1, typemax(UInt64))
+	method.recursion_relation = function (method1, method2, parent_sig, new_sig)
+		return true
+	end
+end
+
 end

test/stage2_fwd.jl

@@ -14,9 +14,7 @@ module stage2_fwd
 
     self_minus(a) = myminus(a, a)
     let self_minus′′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(self_minus), Float64}, 2)
-        # TODO: The IR for this currently contains Union{Diffractor.TangentBundle{2, Float64, Diffractor.ExplicitTangent{Tuple{Float64, Float64, Float64}}}, Diffractor.TangentBundle{2, Float64, Diffractor.TaylorTangent{Tuple{Float64, Float64}}}}
-        # We should have Diffractor be able to prove uniformity
-        @test_broken isa(self_minus′′, Core.OpaqueClosure{Tuple{Float64}, Float64})
+        @test isa(self_minus′′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus′′(1.0) == 0.
     end
 end