perf: fold Statistics moments in four lanes and make min/max branchless#415
Open
agene0001 wants to merge 2 commits into
Open
perf: fold Statistics moments in four lanes and make min/max branchless#415agene0001 wants to merge 2 commits into
agene0001 wants to merge 2 commits into
Conversation
`OnlineMoments` (added in statrs-dev#394 and wired into `Statistics` in cf11836 for this issue) cured the catastrophic case statrs-dev#376 opened with, but Welford is still poorly conditioned when the data carries a large offset: `mean += delta / n` cannot represent a small increment against a large running mean, so the low bits of every update are dropped. On the dataset from the issue - `1e12 + U(0, 1)`, n = 1e6 - measured against a Neumaier-compensated two-pass reference of 8.336923e-2: before 2.5e-4 relative error after 5e-15 Central moments are invariant under a shift, so accumulating moments of `x - first_observation` keeps every magnitude small. It costs one subtraction per observation - 1107 us vs 1103 us over 1e6 elements, i.e. free - and is better conditioned than plain Welford, which carries the same ~2.5e-4 here because it has the same mean-update problem. Also adds `OnlineMoments::merge`, the Chan-Golub-LeVeque pairwise update, which is the parallelisability the issue asks for. It has to reconcile two different offsets, so it is reviewed alongside them. Folding into several accumulators and merging is also slightly better conditioned than one long chain, since each chain accumulates over fewer updates. Refs statrs-dev#376
Two changes to the `Statistics` iterator impls, measured over 1e6 f64:
`mean`/`variance`/`population_variance` folded into a single `OnlineMoments`
chain, which serialises one division per element. Folding into four round-robin
accumulators and merging with `OnlineMoments::merge` keeps four divisions in
flight:
mean 4.40 ms -> 1.11 ms (4.0x)
variance 4.29 ms -> 1.12 ms (3.8x)
This is also slightly *better* conditioned, since each chain accumulates its
rounding over a quarter of the updates: on alternating 1e8/1 data the variance
lands within 0.5 of the exact value against 10 before. It keeps Welford's
overflow robustness - `mean([1.5e308, 1.6e308, 1.7e308])` is 1.6e308, where a
naive sum gives `inf`.
`min`/`max`/`abs_min`/`abs_max` used `if x < acc || x.is_nan()`, whose branch
blocks vectorisation. Using the branchless `f64::min`/`f64::max` with a separate
NaN flag folded in the same pass is 6.0x faster (1.11 ms -> 0.185 ms) and keeps
the documented "any NaN gives NaN" contract, which `f64::min` alone would not
(it returns the non-NaN operand).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two changes to the
Statisticsiterator impls, measured over 1e6 f64 (Apple M3 Pro):Four-lane moment folding.
mean/variance/population_variancefolded a singleOnlineMomentschain, which serialises one division per element. Folding round-robin into four accumulators and merging keeps four divisions in flight:meanvarianceAccuracy is equal or better — each chain accumulates rounding over a quarter of the updates; on alternating 1e8/1 data the variance lands within 0.5 of exact vs 10 before. Welford's overflow robustness is kept:
mean([1.5e308, 1.6e308, 1.7e308])= 1.6e308 where a naive sum givesinf.Branchless min/max.
min/max/abs_min/abs_maxusedif x < acc || x.is_nan(), whose branch blocks vectorisation. Branchlessf64::min/f64::maxwith a separate NaN flag folded in the same pass is 6.0x faster (1.11 ms → 0.185 ms) and preserves the documented "any NaN ⇒ NaN" contract — whichf64::minalone would not, since it returns the non-NaN operand.All existing doctests (NaN propagation, empty input) pass unchanged.
🤖 Generated with Claude Code