Reduce allocations for VarInfo

[mhauru]

Inspired by #1113, and together with @penelopeysm, we tried to squeeze down allocations caused in various parts of VarInfo. This does it to a point where evaluating the LogDensityFunction of a very trivial model causes no allocations what so ever:

julia> module MWE

       using Chairmarks, Distributions, DynamicPPL, LogDensityProblems

       @model function f()
           x ~ Normal()
           y ~ Normal(x)
           return nothing
       end

       function main()
           m = f() | (; y=0.0)
           vi = VarInfo(m)
           ldf = DynamicPPL.LogDensityFunction(m, getlogjoint, vi)
           x = [1.0]
           display(median(@be LogDensityProblems.logdensity(ldf, x)))
       end

       main()
       end
42.772 ns

The same thing on main is:

julia> module MWE

       using Chairmarks, Distributions, DynamicPPL, LogDensityProblems

       @model function f()
           x ~ Normal()
           y ~ Normal(x)
           return nothing
       end

       function main()
           m = f() | (; y=0.0)
           vi = VarInfo(m)
           ldf = DynamicPPL.LogDensityFunction(m, getlogjoint, vi)
           x = [1.0]
           display(median(@be LogDensityProblems.logdensity(ldf, x)))
       end

       main()
       end
128.092 ns (6 allocs: 192 bytes)
Main.MWE

This is not necessarily meant to be merged. We have not thought carefully at all about the effect of @view. This shows that it can be done, and what the important parts are.

[github-actions]

Benchmark Report for Commit 6e633d2

Computer Information

Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

┌───────────────────────┬───────┬─────────────┬───────────────────┬────────┬────────────────┬─────────────────┐
│                 Model │   Dim │  AD Backend │           VarInfo │ Linked │ t(eval)/t(ref) │ t(grad)/t(eval) │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼────────────────┼─────────────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │            4.1 │             2.1 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │          527.1 │            46.8 │
│           Smorgasbord │   201 │ forwarddiff │ simple_namedtuple │   true │          428.0 │            59.7 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │          741.1 │            45.4 │
│           Smorgasbord │   201 │ forwarddiff │       simple_dict │   true │         6710.0 │            26.5 │
│           Smorgasbord │   201 │ forwarddiff │      typed_vector │   true │          817.9 │            39.5 │
│           Smorgasbord │   201 │ forwarddiff │    untyped_vector │   true │          990.0 │            36.3 │
│           Smorgasbord │   201 │ reversediff │             typed │   true │          706.0 │            58.7 │
│           Smorgasbord │   201 │    mooncake │             typed │   true │          653.1 │             5.7 │
│           Smorgasbord │   201 │      enzyme │             typed │   true │            NaN │             NaN │
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │         3095.3 │             4.9 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │          958.8 │             9.3 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │        32359.7 │             4.7 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │         8576.8 │            10.2 │
│               Dynamic │    10 │    mooncake │             typed │   true │          126.0 │            14.2 │
│              Submodel │     1 │    mooncake │             typed │   true │            4.9 │             8.1 │
│                   LDA │    12 │ reversediff │             typed │   true │         1198.8 │             3.6 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴────────────────┴─────────────────┘

[github-actions]

DynamicPPL.jl documentation for PR #1115 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR1115/

[codecov]

Codecov Report

❌ Patch coverage is 94.11765% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 37.48%. Comparing base (1542c56) to head (6e633d2).
⚠️ Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
src/varinfo.jl	93.75%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (1542c56) and HEAD (6e633d2). Click for more details.

HEAD has 18 uploads less than BASE

Flag	BASE (1542c56)	HEAD (6e633d2)
	24	6

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1115       +/-   ##
===========================================
- Coverage   81.23%   37.48%   -43.75%     
===========================================
  Files          40       40               
  Lines        3805     3836       +31     
===========================================
- Hits         3091     1438     -1653     
- Misses        714     2398     +1684     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mhauru]

The only test failure so far (not done though) is an @inferred failure because of SubArray vs Array. This relates to this comment from @torfjelde that this PR hopes to remove:

# TODO(torfjelde): Use `view` instead of `getindex`. Requires addressing type-stability issues though,
# since then we might be returning a `SubArray` rather than an `Array`, which is typically
# what a bijector would result in, even if the input is a view (`SubArray`).
# TODO(torfjelde): An alternative is to implement `view` directly instead.

I don't immediately see why this should be deep problem that couldn't be solved with some function boundaries. @torfjelde, do you have any wise words of warning about there being dragons here?

[penelopeysm]

One potential danger of using views in unflatten is the following aliasing:

julia> using DynamicPPL, Distributions

julia> @model f() = x ~ Normal(); vi = VarInfo(f()); xs = [0.0]
1-element Vector{Float64}:
 0.0

julia> vi = DynamicPPL.unflatten(vi, xs);

julia> vi[@varname(x)] = 1.0  # alternatively: vi = last(DynamicPPL.init!!(f(), vi, InitFromParams((; x = 1.0))))
1.0

julia> xs
1-element Vector{Float64}:
 1.0

I'm fairly sure the aliasing doesn't matter as long as xs, unflatten, and vi are all encapsulated within a LogDensityFunction.

Currently less sure about what happens if unflatten is used outside of LogDensityFunction, or whether it might have knock-on effects in Turing since we abuse unflatten a lot there. For example, if we wrap a parameter vector in an MCMC state, but also unflatten that same vector into a VarInfo and then mutate it, the state will get changed, unless we're careful to copy the vector (which defeats the purpose of the view).

On the other hand I think that using the view in getindex_internal is safe.