split function based on taylor_or_bust

Author: oxinabox

src/codegen/forward_demand.jl

@@ -290,10 +290,10 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 end
                 replace_call!(ir, SSAValue(ssa), Expr(:call, ∂☆{order, eras_mode}(), newargs...))
             elseif isexpr(stmt, :call) || isexpr(stmt, :new)
-                newargs = map(stmt.args) do @nospecialize arg
+                newargs = map(stmt.args) do @nospecialize argq
                     maparg(arg, SSAValue(ssa), order)
                 end
-                f = isexpr(stmt, :call) ? ∂☆{order, eras_mode}() : ∂☆new{order, eras_mode}()
+                f = isexpr(stmt, :call) ? ∂☆{order, eras_mode}() : ∂☆new{order}()
                 replace_call!(ir, SSAValue(ssa), Expr(:call, f, newargs...))
             elseif isa(stmt, PiNode)
                 # TODO: New PiNode that discriminates based on primal?

src/stage1/forward.jl

@@ -74,48 +74,63 @@ function taylor_failure_values(r::TaylorBundle{<:Any, Tuple{Any,Any}}, fail_orde
     return partial(r, i+1)[1], partial(r, i)[2]
 end
 
-function shuffle_up(r::TaylorBundle{1, Tuple{B1,B2}}, ::Val{taylor_or_bust}) where {B1,B2, taylor_or_bust}
-    z₀ = primal(r)[1]
-    z₁ = partial(r, 1)[1]
-    z₂ = primal(r)[2]
-    z₁₂ = partial(r, 1)[2]
-
-    taylor_fail_order = find_taylor_incompatibility(r)
-    if  taylor_fail_order < 0
-        return TaylorBundle{2}(z₀, (z₁, z₁₂))
-    elseif taylor_or_bust
-        @assert taylor_fail_order == 0  # can't be higher
-        throw(TaylorRequired(taylor_fail_order, z₁, z₂))
-    else
-        return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂))
+for taylor_or_bust in (false, true)
+    @eval function shuffle_up(r::TaylorBundle{1, Tuple{B1,B2}}, ::Val{$taylor_or_bust}) where {B1,B2}
+        z₀ = primal(r)[1]
+        z₁ = partial(r, 1)[1]
+        z₂ = primal(r)[2]
+        z₁₂ = partial(r, 1)[2]
+
+        taylor_fail_order = find_taylor_incompatibility(r)
+        if  taylor_fail_order < 0
+            return TaylorBundle{2}(z₀, (z₁, z₁₂))
+        else
+            $(
+                if taylor_or_bust
+                    quote
+                        @assert taylor_fail_order == 0  # can't be higher
+                        throw(TaylorRequired(taylor_fail_order, z₁, z₂))
+                    end
+                else
+                    :(return ExplicitTangentBundle{2}(z₀, (z₁, z₂, z₁₂)))
+                end
+            )
+        end
     end
-end
 
-function shuffle_up(r::TaylorBundle{N, Tuple{B1,B2}}, ::Val{taylor_or_bust}) where {N, B1,B2, taylor_or_bust}
-    the_primal = primal(r)[1]
-    taylor_fail_order = find_taylor_incompatibility(r)
-    if taylor_fail_order(r) < 0
-        the_partials = ntuple(N+1) do i
-            if i <= N
-                partial(r, i)[1]  # == `partial(r,i-1)[2]` (except first which is primal(r)[2])
-            else  # ii = N+1
-                partial(r, i-1)[2]
+    @eval function shuffle_up(r::TaylorBundle{N, Tuple{B1,B2}}, ::Val{$taylor_or_bust}) where {N, B1,B2}
+        the_primal = primal(r)[1]
+        taylor_fail_order = find_taylor_incompatibility(r)
+        if taylor_fail_order(r) < 0
+            the_partials = ntuple(N+1) do i
+                if i <= N
+                    partial(r, i)[1]  # == `partial(r,i-1)[2]` (except first which is primal(r)[2])
+                else  # ii = N+1
+                    partial(r, i-1)[2]
+                end
             end
+            return TaylorBundle{N+1}(the_primal, the_partials)
+        else
+            $(
+                if taylor_or_bust
+                    quote
+                        @assert taylor_fail_order < N
+                        throw(TaylorRequired(taylor_fail_order, taylor_failure_values(r, taylor_fail_order)...))
+                    end
+                else
+                    quote
+                        #XXX: am dubious of the correctness of this
+                        a_partials = ntuple(i->partial(r, i)[1], N)
+                        b_partials = ntuple(i->partial(r, i)[2], N)
+                        the_partials = (a_partials..., primal_b, b_partials...)
+                        return TangentBundle{N+1}(the_primal, the_partials)
+                    end
+                end
+            )
         end
-        return TaylorBundle{N+1}(the_primal, the_partials)
-    elseif taylor_or_bust
-        @assert taylor_fail_order < N
-        throw(TaylorRequired(taylor_fail_order, taylor_failure_values(r, taylor_fail_order)...))
-    else
-        #XXX: am dubious of the correctness of this
-        a_partials = ntuple(i->partial(r, i)[1], N)
-        b_partials = ntuple(i->partial(r, i)[2], N)
-        the_partials = (a_partials..., primal_b, b_partials...)
-        return TangentBundle{N+1}(the_primal, the_partials)
     end
 end
 
-
 function shuffle_up(r::UniformBundle{N, B, U}, _::Val) where {N, B, U}
     (a, b) = primal(r)
     if r.tangent.val === b

test/forward_diff_no_inf.jl

@@ -100,14 +100,14 @@ module forward_diff_no_inf
         #@assert ir[SSAValue(5)][:inst].args[1] == Diffractor.∂☆{1, eras_mode}()
         #@assert ir[SSAValue(5)][:inst].args[2].primal == *
         ir.argtypes[2:end] .= Float64
-        @assert infer_ir!(ir) == Float64
-        
-        Diffractor.forward_diff_no_inf!(ir, [SSAValue(6)] .=> 1; transform! = identity_transform!, eras_mode=eras_mode)
-        ir = CC.compact!(ir)
-        CC.verify_ir(ir)
         infer_ir!(ir)
 
+        Diffractor.forward_diff_no_inf!(ir, [SSAValue(3)] .=> 1; transform! = identity_transform!, eras_mode=eras_mode)
+        # TODO actually test things here.
 
 
+        ir = CC.compact!(ir)
+        CC.verify_ir(ir)
+        infer_ir!(ir)        
     end
 end