Optimise performance for identity VarNames

Author: penelopeysm

src/fastldf.jl

@@ -12,21 +12,35 @@ struct RangeAndLinked
     is_linked::Bool
 end
 
-struct FastLDFContext{T<:AbstractVector{<:Real}} <: AbstractContext
+struct FastLDFContext{N<:NamedTuple,T<:AbstractVector{<:Real}} <: AbstractContext
+    # The ranges of identity VarNames are stored in a NamedTuple for performance
+    # reasons. For just plain evaluation this doesn't make _that_ much of a 
+    # difference (maybe 1.5x), but when doing AD with Mooncake this makes a HUGE
+    # difference (around 4x). Of course, the exact numbers depend on the model.
+    iden_varname_ranges::N
+    # This Dict stores the ranges for all other VarNames
     varname_ranges::Dict{VarName,RangeAndLinked}
+    # The full parameter vector which we index into to get variable values
     params::T
 end
 DynamicPPL.NodeTrait(::FastLDFContext) = IsLeaf()
+function get_range_and_linked(
+    ctx::FastLDFContext, ::VarName{sym,typeof(identity)}
+) where {sym}
+    return ctx.iden_varname_ranges[sym]
+end
+function get_range_and_linked(ctx::FastLDFContext, vn::VarName)
+    return ctx.varname_ranges[vn]
+end
 
 function tilde_assume!!(
     ctx::FastLDFContext, right::Distribution, vn::VarName, vi::OnlyAccsVarInfo
 )
     # Don't need to read the data from the varinfo at all since it's
     # all inside the context.
-    range_and_linked = ctx.varname_ranges[vn]
+    range_and_linked = get_range_and_linked(ctx, vn)
     y = @view ctx.params[range_and_linked.range]
-    is_linked = range_and_linked.is_linked
-    f = if is_linked
+    f = if range_and_linked.is_linked
         from_linked_vec_transform(right)
     else
         from_vec_transform(right)
@@ -51,11 +65,14 @@ end
 struct FastLDF{
     M<:Model,
     F<:Function,
+    N<:NamedTuple,
     AD<:Union{ADTypes.AbstractADType,Nothing},
     ADP<:Union{Nothing,DI.GradientPrep},
 }
     _model::M
     _getlogdensity::F
+    # See FastLDFContext for explanation of these two fields
+    _iden_varname_ranges::N
     _varname_ranges::Dict{VarName,RangeAndLinked}
     _adtype::AD
     _adprep::ADP
@@ -70,6 +87,7 @@ struct FastLDF{
     ) where {syms}
         # Figure out which variable corresponds to which index, and
         # which variables are linked.
+        all_iden_ranges = NamedTuple()
         all_ranges = Dict{VarName,RangeAndLinked}()
         offset = 1
         for sym in syms
@@ -78,7 +96,16 @@ struct FastLDF{
                 len = length(md.ranges[idx])
                 is_linked = md.is_transformed[idx]
                 range = offset:(offset + len - 1)
-                all_ranges[vn] = RangeAndLinked(range, is_linked)
+                if AbstractPPL.getoptic(vn) === identity
+                    all_iden_ranges = merge(
+                        all_iden_ranges,
+                        NamedTuple(
+                            AbstractPPL.getsym(vn) => RangeAndLinked(range, is_linked)
+                        ),
+                    )
+                else
+                    all_ranges[vn] = RangeAndLinked(range, is_linked)
+                end
                 offset += len
             end
         end
@@ -90,23 +117,32 @@ struct FastLDF{
             adtype = tweak_adtype(adtype, model, varinfo)
             x = [val for val in varinfo[:]]
             DI.prepare_gradient(
-                FastLogDensityAt(model, getlogdensity, all_ranges), adtype, x
+                FastLogDensityAt(model, getlogdensity, all_iden_ranges, all_ranges),
+                adtype,
+                x,
             )
         end
 
-        return new{typeof(model),typeof(getlogdensity),typeof(adtype),typeof(prep)}(
-            model, getlogdensity, all_ranges, adtype, prep
+        return new{
+            typeof(model),
+            typeof(getlogdensity),
+            typeof(all_iden_ranges),
+            typeof(adtype),
+            typeof(prep),
+        }(
+            model, getlogdensity, all_iden_ranges, all_ranges, adtype, prep
         )
     end
 end
 
-struct FastLogDensityAt{M<:Model,F<:Function}
+struct FastLogDensityAt{M<:Model,F<:Function,N<:NamedTuple}
     _model::M
     _getlogdensity::F
+    _iden_varname_ranges::N
     _varname_ranges::Dict{VarName,RangeAndLinked}
 end
 function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
-    ctx = FastLDFContext(f._varname_ranges, params)
+    ctx = FastLDFContext(f._iden_varname_ranges, f._varname_ranges, params)
     model = DynamicPPL.setleafcontext(f._model, ctx)
     # This can obviously also be optimised for the case where not
     # all accumulators are needed.
@@ -118,14 +154,23 @@ function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
 end
 
 function LogDensityProblems.logdensity(fldf::FastLDF, params::AbstractVector{<:Real})
-    return FastLogDensityAt(fldf._model, fldf._getlogdensity, fldf._varname_ranges)(params)
+    return FastLogDensityAt(
+        fldf._model, fldf._getlogdensity, fldf._iden_varname_ranges, fldf._varname_ranges
+    )(
+        params
+    )
 end
 
 function LogDensityProblems.logdensity_and_gradient(
     fldf::FastLDF, params::AbstractVector{<:Real}
 )
     return DI.value_and_gradient(
-        FastLogDensityAt(fldf._model, fldf._getlogdensity, fldf._varname_ranges),
+        FastLogDensityAt(
+            fldf._model,
+            fldf._getlogdensity,
+            fldf._iden_varname_ranges,
+            fldf._varname_ranges,
+        ),
         fldf._adprep,
         fldf._adtype,
         params,