Improve Small-Union Performance

src/codual.jl

@@ -33,6 +33,7 @@ Shorthand for `CoDual{P, tangent_type(P}}` when `P` is concrete, equal to `CoDua
 """
 function codual_type(::Type{P}) where {P}
     P == DataType && return CoDual
+    P isa Union && return Union{codual_type(P.a), codual_type(P.b)}
     return isconcretetype(P) ? CoDual{P, tangent_type(P)} : CoDual
 end
 

src/interpreter/reverse_mode_ad.jl

@@ -11,14 +11,20 @@ const BwdsInst = Core.OpaqueClosure{Tuple{Int}, Int}
 
 const RuleSlot{V} = Union{SlotRef{V}, ConstSlot{V}} where {V<:Tuple{CoDual, Ref}}
 
-primal_type(::AbstractSlot{<:Tuple{<:CoDual{P}, <:Any}}) where {P} = @isdefined(P) ? P : Any
-primal_type(::AbstractSlot{<:Tuple{<:CoDual, <:Any}}) = Any
+__primal_type(::Type{<:Tuple{<:CoDual{P}, <:Any}}) where {P} = @isdefined(P) ? P : Any
+function __primal_type(::Type{P}) where {P<:Tuple{<:CoDual, <:Any}}
+    P isa Union && return Union{__primal_type(P.a), __primal_type(P.b)}
+    return Any
+end
 
-function make_rule_slot(::SlotRef{P}, ::Any) where {P}
-    return SlotRef{Tuple{codual_type(P), tangent_ref_type_ub(P)}}()
+primal_type(::AbstractSlot{P}) where {P} = __primal_type(P)
+
+function rule_slot_type(::Type{P}) where {P}
+    return Tuple{codual_type(P), tangent_ref_type_ub(P)}
 end
-function make_rule_slot(::SlotRef{P}, ::PhiNode) where {P}
-    return SlotRef{Tuple{codual_type(P), tangent_ref_type_ub(P)}}()
+
+function make_rule_slot(::SlotRef{P}, ::Any) where {P}
+    return SlotRef{rule_slot_type(P)}()
 end
 function make_rule_slot(x::ConstSlot{P}, ::Any) where {P}
     cd = uninit_codual(x[])
@@ -112,17 +118,18 @@ function build_coinsts(
     tangent_stack_stack = make_tangent_ref_stack(tangent_ref_type_ub(primal_type(val)))
 
     make_fwds(v) = R(primal(v), tangent(v))
-    fwds_inst = @opaque function (p::Int)
+    function fwds_run()
         v, tangent_stack = val[]
         push!(my_tangent_stack, tangent(v))
         push!(tangent_stack_stack, tangent_stack)
         ret[] = (make_fwds(v), top_ref(my_tangent_stack))
         return next_blk
     end
-    bwds_inst = @opaque function (j::Int)
+    fwds_inst = @opaque (p::Int) -> fwds_run()
+    function bwds_run()
         increment_ref!(pop!(tangent_stack_stack), pop!(my_tangent_stack))
-        return j
     end
+    bwds_inst = @opaque (j::Int) -> (bwds_run(); return j)
     return fwds_inst::FwdsInst, bwds_inst::BwdsInst
 end
 

src/tangents.jl

@@ -212,8 +212,7 @@ end
     ))
 
     # If the type is a Union, then take the union type of its arguments.
-    # P isa Union && return Union{tangent_type(P.a), tangent_type(P.b)}
-    P isa Union && return Any
+    P isa Union && return Union{tangent_type(P.a), tangent_type(P.b)}
 
     # If the type is itself abstract, it's tangent could be anything.
     # The same goes for if the type has any undetermined type parameters.

src/test_utils.jl

@@ -1412,6 +1412,15 @@ end
     return @inbounds x[1] + test_recursive_sum(x[2:end])
 end
 
+# Copied over from https://github.com/TuringLang/Turing.jl/issues/1140
+function _sum(x)
+    z = 0
+    for i in eachindex(x)
+        z += x[i]
+    end
+    return z
+end
+
 function generate_test_functions()
     return Any[
         (false, :allocs, nothing, const_tester),
@@ -1587,6 +1596,7 @@ function generate_test_functions()
             randn(sr(2), 700, 500),
             randn(sr(3), 300, 700),
         ),
+        (false, :none, nothing, _sum, randn(1024)),
     ]
 end
 

test/codual.jl

@@ -8,4 +8,8 @@
     @test codual_type(Real) == CoDual
     @test codual_type(Any) == CoDual
     @test codual_type(Type{UnitRange{Int}}) == CoDual{Type{UnitRange{Int}}, NoTangent}
+    @test(==(
+        codual_type(Union{Float64, Int}),
+        Union{CoDual{Float64, Float64}, CoDual{Int, NoTangent}},
+    ))
 end