perf(efficient_did): dispatch conditional Omega* ridge solves to a Rust batched-Cholesky kernel (fits 1.4-1.8x)

CHANGELOG.md

@@ -258,6 +258,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   the no-covariates path is byte-identical and `omega_ridge=0` still routes the entire legacy
   path. A `pt_assumption="post"` covariate fit builds no tables at all (all cells are
   just-identified), keeping that path regression-free.
+- **`EfficientDiD` conditional-path ridge solves dispatch to a Rust batched-Cholesky kernel**
+  (follow-up to the table hoisting above; Rust backend only). `_ridge_solve_weights`' batched
+  `np.linalg.solve` — a serial LAPACK LU sweep over the `(units × H × H)` ridged Omega* stacks
+  and the top stage at every scale post-hoisting — now routes to
+  `batched_ridge_chol_solve_ones` in the Rust backend: a hand-rolled unblocked Cholesky (the
+  ridged Omega* is SPD by construction) in a reused per-thread scratch buffer, parallelized
+  over the unit axis with rayon. Non-SPD rows (measured zero on realistic panels) fall back to
+  LU in-kernel, and any non-finite row is recomputed through the exact legacy numpy chain, so
+  the pseudoinverse edge-case semantics are unchanged. Measured (3-rep medians, 20 periods,
+  5 cohorts, 5 covariates, `aggregate="all"`): fit 7.1s → 4.3s at 2k units (1.65x),
+  36.2s → 22.7s at 10k (1.6x), 129s → 90s at 30k (1.4x), the 100k-unit fit 17.8 → 15.9 min at
+  slightly lower peak RSS; survey-weighted 3.0s → 1.7s at 1k (1.8x). The ridge-solve stage
+  itself is 4.7x faster at 10k (17.1s → 3.7s; the kernel alone is ~9.6x — the remainder is
+  shared Python-side prep, tracked in TODO.md). Results on the Rust backend move at
+  floating-point reassociation level only (post-treatment cells ~2e-12 relative, overall ATT
+  ~1e-13); the pure-Python backend is byte-identical, the no-covariates path is untouched, and
+  `omega_ridge=0` never reaches the kernel.
 - **CallawaySantAnna multiplier bootstrap rewritten as a fused, column-tiled scatter-GEMM**
   (EfficientDiD's bootstrap routes through the same kernel). The former loop sliced the
   `(block × n_units)` weight matrix twice per (g,t) cell per weight block — at a 40-period,

TODO.md

@@ -46,7 +46,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
-| `EfficientDiD` conditional path: post-hoisting (2026-07, 100k fit ~18 min) the top stage at every scale is `_ridge_solve_weights` — batched `np.linalg.solve` LU over `(n × H × H)` stacks (17.2s of a 35.7s fit at 10k units, ~170s at 100k; linear in n). The matrices are SPD after the ridge, so a batched Cholesky (potrf ~1/2 the LU flops) in the Rust backend — which could also thread across the stack — is the natural lever; numpy has no batched triangular solve, so pure-Python gains are limited. | `efficient_did_covariates.py::_ridge_solve_weights`, `rust/src/` | CS-scaling | Mid | Low |
+| `EfficientDiD` conditional path: post-Cholesky-dispatch (2026-07, 10k fit 22.7s) the ridge-solve stage is 3.7s, of which the Rust kernel is only ~1.6s — the remainder is shared Python prep inside `_ridge_solve_weights`: the `zero_mask` full-stack abs scan plus the `omega_stack[rest]` fancy-index copy (~28 GB of extra memory traffic per 10k fit). The copy is avoidable when no rows are zero-masked (`rest.size == n`, the common case) — value-identical on both backends since `om` is never mutated. After that, the largest O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). | `efficient_did_covariates.py::_ridge_solve_weights` | CS-scaling | Quick | Low |
 | `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 |
 | 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 |
 | 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 |

diff_diff/_backend.py

@@ -73,6 +73,15 @@
 except ImportError:
     _rust_demean_map = None
 
+# Batched ridge-regularized SPD solve (EfficientDiD per-unit weights):
+# imported independently for the same mixed-version reason as demean_map.
+try:
+    from diff_diff._rust_backend import (
+        batched_ridge_chol_solve_ones as _rust_batched_ridge_chol_solve,
+    )
+except ImportError:
+    _rust_batched_ridge_chol_solve = None
+
 # Determine final backend based on environment variable and availability
 if _backend_env == "python":
     # Force pure Python mode - disable Rust even if available
@@ -83,6 +92,8 @@
     _rust_compute_robust_vcov = None
     # FE-absorption MAP demeaning kernel
     _rust_demean_map = None
+    # Batched ridge-regularized SPD solve
+    _rust_batched_ridge_chol_solve = None
     # TROP estimator acceleration (local method)
     _rust_unit_distance_matrix = None
     _rust_loocv_grid_search = None
@@ -136,6 +147,8 @@ def rust_backend_info():
     "_rust_compute_robust_vcov",
     # FE-absorption MAP demeaning kernel
     "_rust_demean_map",
+    # Batched ridge-regularized SPD solve (EfficientDiD per-unit weights)
+    "_rust_batched_ridge_chol_solve",
     # TROP estimator acceleration (local method)
     "_rust_unit_distance_matrix",
     "_rust_loocv_grid_search",

diff_diff/efficient_did_covariates.py

@@ -26,6 +26,7 @@
 import numpy as np
 from scipy.spatial.distance import cdist
 
+from diff_diff._backend import HAS_RUST_BACKEND, _rust_batched_ridge_chol_solve
 from diff_diff.linalg import solve_ols
 
 # Default ridge for the Omega* inversion (see ``compute_per_unit_weights`` /
@@ -1629,27 +1630,57 @@ def _ridge_solve_weights(omega_stack: np.ndarray, omega_ridge: float) -> np.ndar
 
     om = omega_stack[rest]
     trace = np.trace(om, axis1=1, axis2=2)
-    scale = np.maximum(trace / H, 0.0)
-    om_ridged = om + (omega_ridge * scale)[:, None, None] * np.eye(H)[None]
+    ridge = omega_ridge * np.maximum(trace / H, 0.0)
+    if (
+        HAS_RUST_BACKEND
+        and _rust_batched_ridge_chol_solve is not None
+        and omega_stack.dtype == np.float64
+    ):
+        # Batched Cholesky in Rust (ridge added in-kernel - no om_ridged
+        # temp). Dispatch is on availability + dtype ONLY: any batch-size
+        # cutoff would make forced-tiny-tile fits cross-algorithm vs the
+        # single-tile fit and break the tile-invariance twins. A non-finite
+        # row signals not-SPD (kernel LU fallback / NaN-poisoned exact
+        # singularity) - recompute exactly those rows via the legacy chain
+        # so their semantics (incl. the pinv arm) match the numpy path.
+        num = _rust_batched_ridge_chol_solve(om, ridge)
+        bad = np.flatnonzero(~np.isfinite(num).all(axis=1))
+        if bad.size:
+            num[bad] = _ridge_solve_numpy(om[bad], ridge[bad])
+    else:
+        num = _ridge_solve_numpy(om, ridge)
+    den = num @ ones
+    ok = np.abs(den) >= 1e-15
+    solved = np.full((rest.size, H), 1.0 / H)
+    solved[ok] = num[ok] / den[ok, None]
+    weights[rest] = solved
+    return weights
+
+
+def _ridge_solve_numpy(om: np.ndarray, ridge: np.ndarray) -> np.ndarray:
+    """Legacy numpy solve chain for ``(om_i + ridge_i * I) x = 1``.
+
+    Batched LU solve; on an exactly-singular batch, per-unit solve with a
+    minimum-norm pseudoinverse backstop (unreachable for PSD Omega* with
+    lam > 0, kept so one pathological unit cannot fail the batch). This is
+    both the pure-Python path and the per-row recompute target for rows the
+    Rust kernel flags as non-finite.
+    """
+    m_rows, H, _ = om.shape
+    om_ridged = om + ridge[:, None, None] * np.eye(H)[None]
     try:
         # gufunc solve treats a 2-D rhs as a matrix; a stack of vectors needs
-        # shape (n, H, 1)
-        num = np.linalg.solve(om_ridged, np.ones((rest.size, H, 1)))[..., 0]
+        # shape (m, H, 1)
+        return np.linalg.solve(om_ridged, np.ones((m_rows, H, 1)))[..., 0]
     except np.linalg.LinAlgError:
-        # Unreachable for PSD Omega* with lam > 0; kept as a per-unit
-        # minimum-norm fallback so one pathological unit cannot fail the batch.
-        num = np.empty((rest.size, H))
-        for m in range(rest.size):
+        ones = np.ones(H)
+        num = np.empty((m_rows, H))
+        for m in range(m_rows):
             try:
                 num[m] = np.linalg.solve(om_ridged[m], ones)
             except np.linalg.LinAlgError:
                 num[m] = np.linalg.pinv(om_ridged[m]) @ ones
-    den = num @ ones
-    ok = np.abs(den) >= 1e-15
-    solved = np.full((rest.size, H), 1.0 / H)
-    solved[ok] = num[ok] / den[ok, None]
-    weights[rest] = solved
-    return weights
+        return num
 
 
 def compute_per_unit_weights(