Cache transition CDFs on MarkovChain with thread-safe lazy initialization

[oyamad]

#395 originally cached the transition CDFs used by the simulation routines on the MarkovChain instance (lazily, on first use). The cache was removed before merge because the naive lazy initialization is not thread safe; this issue records the motivation and a concrete design for reintroducing it.

Motivation

With per-call CDF construction (the state after #395), repeated short simulations of a large dense chain pay O(n²) setup per call: simulate/dense_n1000_ts100 is ~930 μs per call versus 3.6 μs with a warm cache (numbers from the initial #395 revision; ratios are what matter). QuantEcon.py caches, in the cdfs/cdfs1d properties of quantecon/markov/core.py. simulate! amortizes construction across the columns of one call; a cache would amortize across calls.

Why the naive version was removed

The removed implementation (see b311f7f, reverted in #395) did:

s = getfield(mc, :sampler)
if s === nothing
    s = _sampler_for(mc.p)
    setfield!(mc, :sampler, s)
end

The first simulate call therefore mutates the chain, so the natural threaded pattern — Threads.@threads over independent sample paths from a shared mc — races on the field. The race is benign on current hardware (both threads build identical samplers; the field is a single boxed pointer), but it is a data race and hence undefined behavior under Julia's memory model, silently so. QuantEcon.py has the same unlocked pattern but is protected by the GIL, by numba kernels that hold the GIL, and by process-based parallelism culture; none of these apply to Julia, where threads are the primary parallelism story.

Proposed design: atomic lazy initialization

Julia ≥ 1.7 (compat is 1.10) supports atomic fields:

mutable struct MarkovChain{T, TM, TV, TS}
    p::TM
    state_values::TV
    @atomic sampler::Union{Nothing,TS}
    ...
end

function _get_sampler(mc::MarkovChain)
    s = @atomic :acquire mc.sampler
    s === nothing || return s
    s_new = _sampler_for(mc.p)
    old, ok = @atomicreplace mc.sampler nothing => s_new
    return ok ? s_new : something(old)
end

The something(old) narrows the Union{Nothing,TS} type of the CAS witness: a failed CAS guarantees it saw a non-nothing value, but inference does not know that, and without the narrowing _get_sampler would infer as a Union and fail the @inferred tests.

Concurrent first-callers may both build, but the compare-and-swap installs exactly one sampler and every caller uses a fully constructed one. The duplicate O(n²) construction under a cold race is accepted deliberately: it happens at most once per chain and is small next to the simulation work that motivates threading. (If it ever matters, the alternative is an atomic fast path with a lock taken only on the cold path.) The warm path costs one acquire load (free on x86, nearly free on ARM). Invalidation on mc.p assignment (a Base.setproperty! override, as in the removed implementation) must also write the field atomically.

Guarantee to provide and document: concurrent read-only operations on a shared chain are safe, cold or warm. Concurrently mutating mc.p while another thread simulates remains the user's responsibility, per the usual Julia convention. In-place mutation (mc.p .= new_P) is invisible to any setproperty! hook, so the documented contract must be that in-place modification requires a subsequent assignment (mc.p = mc.p suffices) to invalidate the cache — as in the docstring of the removed implementation.

Alternatives considered

• Eager construction in the constructor: no concurrency machinery at all, but every chain pays the O(n²) memory and construction cost, including uses that never simulate.
• ReentrantLock-guarded lazy initialization: correct without atomics, at the cost of a lock/unlock per simulate call (tens of ns; the fast path cannot skip the lock without an atomic read).
• Status quo: per-call construction; only the repeated-short-simulation pattern suffers, and simulate! batching is a workaround.

Notes from the removed implementation

• A 4th type parameter TS computed from TM keeps the field access and the simulation entry points type stable (@inferred-tested).
• _sampler_type predicted the dense CDF eltype with Base.promote_op(Base.add_sum, T, T) while _sampler_for used whatever cumsum returns; these are Base internals and can drift. Better: have _sampler_for construct the declared type explicitly so they agree by construction.
• The setproperty! override should convert first and invalidate only after a successful assignment.
• Add a threaded stress test (e.g. first-use from many threads) and a regression test that assignment to mc.p invalidates the cache.

🤖 Generated with Claude Code (Claude Fable 5)