Skip to content

[upstream-candidate] Hoist build_tree scratch vectors into depth-indexed members (base_nuts) - #1

Open
sims1253 wants to merge 1 commit into
developfrom
scratch-hoist-base-nuts
Open

[upstream-candidate] Hoist build_tree scratch vectors into depth-indexed members (base_nuts)#1
sims1253 wants to merge 1 commit into
developfrom
scratch-hoist-base-nuts

Conversation

@sims1253

@sims1253 sims1253 commented Aug 22, 2026

Copy link
Copy Markdown
Owner

Hoists the per-recursion-level locals in base_nuts::build_tree (p_init_end, p_sharp_init_end, rho_init, p_final_beg, p_sharp_final_beg, rho_final, rho_subtree, z_propose_final) plus transition()'s rho_extended into member scratch buffers, sized lazily on first transition (num_params x max_depth). Removes ~630 heap allocations per transition at tree depth 6 on hierarchical models.

Recursion safety: a node at depth d only touches slot d; children use slot d-1; outputs children write into the parent's slot are read only after the children return. z_propose_final is a per-depth ps_point stack — a single shared ps_point would be overwritten by the recursive call that receives it as z_propose.

Untouched: leaf z_propose assignment, transition-scope vectors, compute_criterion, integrator/hamiltonian internals, RNG call order. No arithmetic reordering.

Evidence (cmdstan 2.39.0; models blr, pilots, lsat_model, arma11, eight_schools_noncentered)

  • Bit-identity: 3 models x 2 seeds x 4 chains vs stock: 24/24 CSVs byte-identical (modulo elapsed-time/file-path comments).
  • callgrind (pilots, 40 warmup + 40 draws): memcpy/alloc instruction share 9.9% -> 6.7%; total Ir -6%.
  • Wall-clock (warmup+sampling, 1000+1000, 4 chains, medians of 3 reps): geomean ratio patched/stock 0.931 (pilots 0.928, arma11 0.966, blr 0.888, eight_schools_noncentered 0.943).

transition()-time heap churn measured ~630 allocations per transition at
tree depth 6 on hierarchical models; each build_tree recursion level
allocates p_init_end, p_sharp_init_end, rho_init, p_final_beg,
p_sharp_final_beg, rho_final, rho_subtree and copies a ps_point
z_propose_final. This hoists all of them into depth-indexed member
scratch buffers sized lazily on the first transition (num_params x
max_depth), plus a single member for transition()'s rho_extended.

Recursion safety: a node at depth d only touches slot d; its children
use slot d-1, and outputs written by children into the parent's slot are
read only after the children return. z_propose_final must be a
per-depth ps_point stack - a single shared ps_point would be overwritten
by the recursive call that receives it as z_propose.

Pure memory-management change: no arithmetic reordering, no RNG call
changes, transition-scope vectors, the leaf z_propose assignment,
compute_criterion and integrator/hamiltonian internals untouched.

Verified (cmdstan 2.39.0 build with models blr, pilots, lsat_model,
arma11, eight_schools_noncentered):
- bit-identical CSVs vs stock: 3 models x 2 seeds x 4 chains, 24/24
- callgrind (pilots, 40+40): memcpy/alloc instruction share 9.9% -> 6.7%,
  total Ir -6%
- wall-clock (warmup+sampling, 1000+1000, 4 chains, medians of 3 reps):
  geomean ratio patched/stock 0.931 (pilots 0.928, arma11 0.966,
  blr 0.888, eight_schools_noncentered 0.943)
@sims1253
sims1253 force-pushed the scratch-hoist-base-nuts branch from 7fc7f7e to 092f040 Compare August 23, 2026 18:44
@sims1253 sims1253 changed the title Hoist build_tree scratch vectors into depth-indexed members (base_nuts) [internal] Hoist build_tree scratch vectors into depth-indexed members (base_nuts) Aug 25, 2026
@sims1253 sims1253 changed the title [internal] Hoist build_tree scratch vectors into depth-indexed members (base_nuts) [upstream-candidate] Hoist build_tree scratch vectors into depth-indexed members (base_nuts) Aug 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant