Skip to content

Commit abf21bd

Browse files
igerberigerber
authored
perf(linalg): slim solve_ols marshalling on both backends - 15.3 -> 7.8 GB rust-side high-water, CV-clear speed wins, outputs preserved (#607)
1 parent c3ea543 commit abf21bd

15 files changed

Lines changed: 847 additions & 309 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
3333
attribution in `docs/performance-plan.md`).
3434

3535
### Changed
36+
- **`solve_ols` marshalling slimmed on both backends** (memory + speed; outputs preserved).
37+
Rust: column norms computed from the borrowed numpy view, equilibration fused into the
38+
single owned faer copy, `U'y` computed directly off the faer SVD factor (the old path
39+
materialized a full ndarray copy of U for one k-vector dot), and factors dropped before
40+
the fitted/residual/vcov stage - rust-side allocator high-water on a 2.4M x 130 clustered
41+
solve falls from 15.32 GB to 7.81 GB (measured with the new feature-gated `alloc-profile`
42+
counting allocator; never compiled into wheels). Python: rank detection uses
43+
`qr(mode="r")` (same dgeqp3 R and pivot, skips forming the discarded Q) and
44+
`_equilibrated_lstsq` materializes its scaled temporary F-order with `overwrite_a=True`
45+
so gelsd consumes it in place - python-backend estimates are bit-identical (benchmark
46+
identity deltas exactly 0.0). The removed copies also cost time: frozen-code benchmark
47+
arms measured CV-clear improvements incl. firm-panel rust fits -21% (25.8 -> 20.3 s),
48+
county event studies -12%/-13% (rust/python), with no scenario regressing. Rust outputs
49+
move only at the cross-backend parity tolerances (REGISTRY-documented; never
50+
bit-identical across backends).
3651
- **FE-absorption demeaning rewritten: factorize-once + `np.bincount` method of alternating
3752
projections** (`demean_by_groups` / `within_transform`). Each absorbed dimension is factorized
3853
once and group means are formed via `np.bincount` instead of re-hashing the group keys with

‎TODO.md‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,13 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
4646

4747
| Issue | Location | Origin | Effort | Priority |
4848
|-------|----------|--------|--------|----------|
49-
| Rust-backend `solve_ols` peak memory on wide designs: the faer path holds an owned copy of the stacked design plus SVD workspace, dominating fit-level peak RSS on wide absorbed fits (firm_churn-class: ~19-21 GB rust-backend vs ~13.4 GB numpy-backend for the same fit; measured PR-D 2026-07, corrected attribution in docs/performance-plan.md - the demean kernel's transients were NOT the dominator, and opt-in chunked dispatch caps them anyway). Candidates: borrow-view instead of owned copy into faer, workspace reuse, or routing wide designs to the numpy gelsd path. Interim escape hatch: DIFF_DIFF_BACKEND=python. | `rust/src/linalg.rs`, `diff_diff/linalg.py::solve_ols` | PR-D probe | Heavy | Medium |
5049
| `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium |
5150
| CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low |
5251
| Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low |
5352
| `SyntheticControl` retains a full `_SyntheticControlFitSnapshot` (pivoted panels) on EVERY fit for the opt-in `in_space_placebo()`, so callers who never run the placebo pay `O(units × periods × predictor-vars)` memory. Store a compact array/index representation, or build the snapshot lazily on first placebo call. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Mid | Low |
5453
| Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low |
5554
| `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low |
5655
| `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low |
57-
| Rust faer SVD ndarray-to-faer conversion overhead (minimal vs SVD cost). | `rust/src/linalg.rs:67` | #115 | Quick | Low |
5856
| Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low |
5957
| `ImputationDiD._iterative_demean` (`imputation.py:1064`) and its near-identical twin `TwoStageDiD._iterative_demean` (`two_stage.py:2201`) rebuild a `pd.Series`+`groupby().transform()` every alternating-projection iteration. Precompute the unit/time group codes once and use `np.bincount` for the per-iteration group means. **Not bit-identical** on pandas 3.0 (`np.bincount` is naive accumulation; pandas' `group_mean` is Kahan-compensated → ~5.8e-11, the same order as the demean's `tol=1e-10`), so this needs a tolerance/REGISTRY-Note treatment + golden re-validation. Modest peak effect (the per-iteration Series are transient), analytical-path-only (the bootstrap reuses #562's cached projection). Optimize both twins together (cross-surface-twins). | `imputation.py`, `two_stage.py` | PR-C deferral | Mid | Low |
6058

@@ -124,6 +122,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope.
124122
| Issue | Location | PR | Priority |
125123
|-------|----------|----|----------|
126124
| `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low |
125+
| `solve_ols` residual memory floor is inherent to SVD-based lstsq after the PR-E marshalling slim (rust-side allocator high-water 15.32 -> 7.81 GB on the 2.4M x 130 clustered solve; remaining = one fused equilibrated input copy + thin-SVD U transient + vcov scores block + LAPACK gelsd internals on the python path). Further reduction requires the tall-skinny-QR algorithm change - same trade-off as the parked QR-reuse row below (library-wide last-digit perturbation + loses the second independent collinearity detector); the two land together if ever reopened. Diagnosis + measurements in docs/performance-plan.md. | `rust/src/linalg.rs::solve_ols`, `diff_diff/linalg.py` | PR-E | Low |
127126
| SunAbraham-family fits are SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel; the pivoted QR alone is ~20% of the whole fit (measured attribution in docs/performance-plan.md). Candidate fix is QR-reuse (one factorization for rank detection + solve). **Parked per the 2026-07 correctness-first decision**: the change perturbs every coefficient in the library at the last digits (golden/parity/bit-identity recapture) and removes the SVD's independent singular-value truncation - a second, different near-collinearity detector that is deliberate defense-in-depth for the no-silent-failures contract (the FE-span guard episode showed borderline cases are real). Re-open only on practitioner demand for wide event studies; any implementation gates on fixest/R end-to-end parity, not just internal consistency. | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Medium |
128127
| dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low |
129128
| `HeterogeneousAdoptionDiD` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium |

‎benchmarks/speed_review/baselines/fe_absorption_after.json‎

Lines changed: 39 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,14 @@
2525
"checksum": 6315.071127371084,
2626
"att": 0.2803785413990301,
2727
"se": 0.031244864610215972,
28-
"fit_median_s": 2.3882,
29-
"fit_min_s": 2.3505,
30-
"fit_max_s": 2.4948,
31-
"fit_cv": 0.0257,
28+
"fit_median_s": 2.0868,
29+
"fit_min_s": 2.0496,
30+
"fit_max_s": 2.1839,
31+
"fit_cv": 0.0255,
3232
"noisy": false,
3333
"n_timed_fits": 9,
34-
"datagen_s": 0.004,
35-
"peak_rss_mb": 3177.5,
34+
"datagen_s": 0.005,
35+
"peak_rss_mb": 2855.8,
3636
"backend_resolved": "python",
3737
"warnings": []
3838
},
@@ -42,14 +42,14 @@
4242
"checksum": 240997.06223418942,
4343
"att": 0.3144327184497798,
4444
"se": 0.005837281701253927,
45-
"fit_median_s": 49.502,
46-
"fit_min_s": 48.4532,
47-
"fit_max_s": 50.3118,
48-
"fit_cv": 0.0189,
45+
"fit_median_s": 46.1275,
46+
"fit_min_s": 44.4378,
47+
"fit_max_s": 46.8508,
48+
"fit_cv": 0.027,
4949
"noisy": false,
5050
"n_timed_fits": 3,
51-
"datagen_s": 0.07,
52-
"peak_rss_mb": 15900.8,
51+
"datagen_s": 0.072,
52+
"peak_rss_mb": 15072.3,
5353
"backend_resolved": "python",
5454
"warnings": []
5555
},
@@ -59,14 +59,14 @@
5959
"checksum": 76832.22571739106,
6060
"att": 0.251164131495566,
6161
"se": 0.002414938742160613,
62-
"fit_median_s": 0.9834,
63-
"fit_min_s": 0.9825,
64-
"fit_max_s": 0.9845,
65-
"fit_cv": 0.001,
62+
"fit_median_s": 0.9799,
63+
"fit_min_s": 0.9765,
64+
"fit_max_s": 0.9907,
65+
"fit_cv": 0.0076,
6666
"noisy": false,
6767
"n_timed_fits": 3,
68-
"datagen_s": 0.077,
69-
"peak_rss_mb": 976.9,
68+
"datagen_s": 0.076,
69+
"peak_rss_mb": 902.4,
7070
"backend_resolved": "python",
7171
"warnings": []
7272
},
@@ -76,14 +76,14 @@
7676
"checksum": 443696.3834000407,
7777
"att": 0.25194372599916653,
7878
"se": 0.0017886314723251327,
79-
"fit_median_s": 0.973,
80-
"fit_min_s": 0.968,
81-
"fit_max_s": 0.9759,
82-
"fit_cv": 0.0041,
79+
"fit_median_s": 0.9421,
80+
"fit_min_s": 0.9417,
81+
"fit_max_s": 0.9425,
82+
"fit_cv": 0.0004,
8383
"noisy": false,
8484
"n_timed_fits": 3,
85-
"datagen_s": 0.062,
86-
"peak_rss_mb": 1664.8,
85+
"datagen_s": 0.061,
86+
"peak_rss_mb": 1668.0,
8787
"backend_resolved": "python",
8888
"warnings": [
8989
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -96,14 +96,14 @@
9696
"checksum": -176482.6014001448,
9797
"att": 0.24554568144130481,
9898
"se": 0.007662698712656388,
99-
"fit_median_s": 4.0767,
100-
"fit_min_s": 4.0736,
101-
"fit_max_s": 4.0778,
102-
"fit_cv": 0.0005,
99+
"fit_median_s": 4.0394,
100+
"fit_min_s": 4.0254,
101+
"fit_max_s": 4.0462,
102+
"fit_cv": 0.0026,
103103
"noisy": false,
104104
"n_timed_fits": 3,
105105
"datagen_s": 0.107,
106-
"peak_rss_mb": 3589.2,
106+
"peak_rss_mb": 3578.0,
107107
"backend_resolved": "python",
108108
"warnings": [
109109
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -116,14 +116,14 @@
116116
"checksum": 376703.7248141762,
117117
"att": 0.2502082560083505,
118118
"se": 0.004392830594147583,
119-
"fit_median_s": 31.7648,
120-
"fit_min_s": 31.4957,
121-
"fit_max_s": 31.8372,
122-
"fit_cv": 0.0057,
119+
"fit_median_s": 31.656,
120+
"fit_min_s": 31.5351,
121+
"fit_max_s": 31.7636,
122+
"fit_cv": 0.0036,
123123
"noisy": false,
124124
"n_timed_fits": 3,
125125
"datagen_s": 0.067,
126-
"peak_rss_mb": 1840.7,
126+
"peak_rss_mb": 1839.7,
127127
"backend_resolved": "python",
128128
"warnings": [
129129
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -136,14 +136,14 @@
136136
"checksum": -571.6190664591859,
137137
"att": 0.2515087011451766,
138138
"se": 0.028726073333820948,
139-
"fit_median_s": 0.0035,
140-
"fit_min_s": 0.0034,
141-
"fit_max_s": 0.0039,
142-
"fit_cv": 0.0437,
139+
"fit_median_s": 0.0036,
140+
"fit_min_s": 0.0035,
141+
"fit_max_s": 0.0042,
142+
"fit_cv": 0.0487,
143143
"noisy": false,
144144
"n_timed_fits": 21,
145145
"datagen_s": 0.001,
146-
"peak_rss_mb": 145.1,
146+
"peak_rss_mb": 148.1,
147147
"backend_resolved": "python",
148148
"warnings": []
149149
}

‎benchmarks/speed_review/baselines/fe_absorption_rust.json‎

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -23,50 +23,50 @@
2323
"scenario": "county_policy",
2424
"n_obs": 177289,
2525
"checksum": 6315.071127371084,
26-
"att": 0.2803785413990308,
27-
"se": 0.03124486461021617,
28-
"fit_median_s": 1.7567,
29-
"fit_min_s": 1.7285,
30-
"fit_max_s": 1.8617,
31-
"fit_cv": 0.0302,
26+
"att": 0.28037854139903057,
27+
"se": 0.031244864610215826,
28+
"fit_median_s": 1.5536,
29+
"fit_min_s": 1.4852,
30+
"fit_max_s": 1.6031,
31+
"fit_cv": 0.0286,
3232
"noisy": false,
3333
"n_timed_fits": 9,
3434
"datagen_s": 0.005,
35-
"peak_rss_mb": 4237.1,
35+
"peak_rss_mb": 3264.3,
3636
"backend_resolved": "rust",
3737
"warnings": []
3838
},
3939
{
4040
"scenario": "firm_churn",
4141
"n_obs": 2400000,
4242
"checksum": 240997.06223418942,
43-
"att": 0.31443271844971477,
44-
"se": 0.005837281701253889,
45-
"fit_median_s": 25.8387,
46-
"fit_min_s": 24.6121,
47-
"fit_max_s": 26.7088,
48-
"fit_cv": 0.041,
43+
"att": 0.314432718449715,
44+
"se": 0.005837281701253807,
45+
"fit_median_s": 20.341,
46+
"fit_min_s": 19.7481,
47+
"fit_max_s": 20.7773,
48+
"fit_cv": 0.0255,
4949
"noisy": false,
5050
"n_timed_fits": 3,
51-
"datagen_s": 0.069,
52-
"peak_rss_mb": 20963.6,
51+
"datagen_s": 0.07,
52+
"peak_rss_mb": 17101.5,
5353
"backend_resolved": "rust",
5454
"warnings": []
5555
},
5656
{
5757
"scenario": "scanner_twfe",
5858
"n_obs": 3255000,
5959
"checksum": 76832.22571739106,
60-
"att": 0.25116413149501093,
61-
"se": 0.0024149387421606118,
62-
"fit_median_s": 0.6239,
63-
"fit_min_s": 0.6215,
64-
"fit_max_s": 0.6252,
65-
"fit_cv": 0.003,
60+
"att": 0.25116413149499167,
61+
"se": 0.0024149387421606113,
62+
"fit_median_s": 0.6118,
63+
"fit_min_s": 0.6118,
64+
"fit_max_s": 0.6154,
65+
"fit_cv": 0.0034,
6666
"noisy": false,
6767
"n_timed_fits": 3,
6868
"datagen_s": 0.076,
69-
"peak_rss_mb": 1477.0,
69+
"peak_rss_mb": 1327.8,
7070
"backend_resolved": "rust",
7171
"warnings": []
7272
},
@@ -76,14 +76,14 @@
7676
"checksum": 443696.3834000407,
7777
"att": 0.25194372599916653,
7878
"se": 0.0017886314723251327,
79-
"fit_median_s": 0.4839,
80-
"fit_min_s": 0.483,
81-
"fit_max_s": 0.4845,
82-
"fit_cv": 0.0016,
79+
"fit_median_s": 0.4525,
80+
"fit_min_s": 0.4515,
81+
"fit_max_s": 0.4566,
82+
"fit_cv": 0.006,
8383
"noisy": false,
8484
"n_timed_fits": 3,
8585
"datagen_s": 0.062,
86-
"peak_rss_mb": 1744.8,
86+
"peak_rss_mb": 1741.9,
8787
"backend_resolved": "rust",
8888
"warnings": [
8989
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -96,14 +96,14 @@
9696
"checksum": -176482.6014001448,
9797
"att": 0.24554568144130481,
9898
"se": 0.007662698712656388,
99-
"fit_median_s": 2.8354,
100-
"fit_min_s": 2.8354,
101-
"fit_max_s": 2.8357,
102-
"fit_cv": 0.0001,
99+
"fit_median_s": 2.778,
100+
"fit_min_s": 2.7747,
101+
"fit_max_s": 2.7822,
102+
"fit_cv": 0.0014,
103103
"noisy": false,
104104
"n_timed_fits": 3,
105-
"datagen_s": 0.106,
106-
"peak_rss_mb": 3587.3,
105+
"datagen_s": 0.105,
106+
"peak_rss_mb": 3586.4,
107107
"backend_resolved": "rust",
108108
"warnings": [
109109
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -116,14 +116,14 @@
116116
"checksum": 376703.7248141762,
117117
"att": 0.2502082560083505,
118118
"se": 0.004392830594147583,
119-
"fit_median_s": 11.365,
120-
"fit_min_s": 11.3007,
121-
"fit_max_s": 11.4856,
122-
"fit_cv": 0.0082,
119+
"fit_median_s": 11.3811,
120+
"fit_min_s": 11.3491,
121+
"fit_max_s": 11.4257,
122+
"fit_cv": 0.0034,
123123
"noisy": false,
124124
"n_timed_fits": 3,
125-
"datagen_s": 0.066,
126-
"peak_rss_mb": 1841.8,
125+
"datagen_s": 0.067,
126+
"peak_rss_mb": 1842.5,
127127
"backend_resolved": "rust",
128128
"warnings": [
129129
"Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are",
@@ -134,16 +134,16 @@
134134
"scenario": "guard_small",
135135
"n_obs": 20000,
136136
"checksum": -571.6190664591859,
137-
"att": 0.2515087011451809,
138-
"se": 0.028726073333820955,
139-
"fit_median_s": 0.0034,
137+
"att": 0.2515087011451813,
138+
"se": 0.028726073333820944,
139+
"fit_median_s": 0.0035,
140140
"fit_min_s": 0.0032,
141-
"fit_max_s": 0.0036,
142-
"fit_cv": 0.0346,
141+
"fit_max_s": 0.0039,
142+
"fit_cv": 0.0512,
143143
"noisy": false,
144144
"n_timed_fits": 21,
145145
"datagen_s": 0.001,
146-
"peak_rss_mb": 148.0,
146+
"peak_rss_mb": 148.3,
147147
"backend_resolved": "rust",
148148
"warnings": []
149149
}

0 commit comments

Comments
 (0)