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

[igerber]

Summary

• _ridge_solve_weights' batched np.linalg.solve — a serial LAPACK LU sweep over the (units x H x H) ridged Omega* stacks, and the top stage at every scale after perf(efficient_did): cross-cell kcov-table hoisting + per-group W lifecycle (100k: 86 -> 18 min) #629's kcov-table hoisting — now dispatches to a new Rust kernel batched_ridge_chol_solve_ones on the Rust backend: hand-rolled unblocked Cholesky (the ridged Omega* is SPD by construction) in a reused per-thread scratch buffer, rayon-parallel over the unit axis.
• Non-SPD rows (measured zero on realistic panels; the ridge floors relative eigenvalues at ~1e-8) fall back to faer LU in-kernel, with NaN-poisoning on an exactly-zero pivot; any non-finite row is recomputed through the exact legacy numpy chain, so the pseudoinverse edge-case semantics are unchanged. Dispatch is on backend availability + float64 only — no batch-size cutoff, keeping the tile-invariance twins same-algorithm on both sides.
• Measured (3-rep medians, 20 periods / 5 cohorts / 5 covariates, aggregate="all"): 2k 7.1 -> 4.3s (1.65x), 10k 36.2 -> 22.7s (1.60x), 30k 129 -> 90s (1.43x), 100k 17.8 -> 15.9 min, survey 1k 1.8x; ridge stage at 10k 17.1 -> 3.7s (kernel alone 9.6x); maxrss flat or lower.
• Results on the Rust backend move at floating-point reassociation level only (post cells ~2e-12 rel, overall ATT ~1e-13, near-zero placebos <= 1.4e-11 abs); the pure-Python backend is byte-identical (verified via a dedicated arm), the no-covariates path is untouched, and omega_ridge=0 never reaches the kernel.

Methodology references (required if estimator / math changes)

• Method name(s): EfficientDiD (Chen, Sant'Anna & Xie 2025) — conditional Omega* efficient-weight solve
• Paper / source link(s): Chen, Sant'Anna & Xie (2025); docs/methodology/REGISTRY.md Omega* ridge Note (extended with the Rust batched-Cholesky follow-up sentence)
• Any intentional deviations from the source (and why): None — pure implementation change; the ridge system (Omega* + lam * max(trace/H, 0) * I) x = 1 and weight normalization are unchanged, Cholesky vs LU differ only at the condition-bounded floating-point level.

Validation

• Tests added/updated: tests/test_rust_backend.py (new TestBatchedRidgeCholSolve: parity vs numpy on well-/ill-conditioned stacks, NaN/exact-singular rows, batch-split bit-identity, degenerate shapes, strided input, shape validation), tests/test_efficient_did.py (new TestRidgeSolveRustDispatch: kernel-called spy, in-process backend A/B, symbol-None degradation, bad-row-only exact recompute, exact-singular pinv arm, f32 declines dispatch), 5 Rust cargo test units. All existing EfficientDiD locks (dense-reference oracle, tile twins, 1-ulp stability, hand-computed weights, HRS anchors) re-run UNMODIFIED on both backends.
• Backtest / simulation / notebook evidence (if applicable): perf arms in docs/performance-plan.md (BEFORE = main 81d53f5 via pristine worktree); tutorial 15 executes clean via nbmake.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

✅ Looks good. I found no unmitigated P0/P1 methodology, correctness, inference, or edge-case issues in the PR diff.

Executive Summary

• Affected method: EfficientDiD, conditional Omega* efficient-weight ridge solve.
• The new Rust Cholesky dispatch preserves the registry-documented system (Omega* + omega_ridge * max(trace/H, 0) * I) x = 1 and normalizes by 1'x.
• Cholesky vs LU is a valid numerical implementation choice here; the deviation is explicitly documented in docs/methodology/REGISTRY.md:L1257.
• Legacy omega_ridge=0 routing remains separate and unchanged.
• The remaining Python prep overhead is tracked in TODO.md:L49, so it is informational only.

Methodology

Finding: Documented Implementation Change, Not A Defect

• Severity: P3 informational
• Location: diff_diff/efficient_did_covariates.py:L1600-L1657, rust/src/batched_solve.rs:L43-L90, docs/methodology/REGISTRY.md:L1257
• Impact: The PR changes the numerical solver used for the conditional Omega* ridge system on the Rust backend, but it preserves the documented target equation and weight normalization. The registry explicitly labels the Rust batched-Cholesky solve as a pure implementation change, and Cholesky vs LU is an acceptable numerical approach difference under the review rules.
• Concrete fix: None required.

Finding: Legacy Ridge-Off Path Preserved

• Severity: P3 informational
• Location: diff_diff/efficient_did.py:L1048-L1063, diff_diff/efficient_did_covariates.py:L1726-L1756
• Impact: omega_ridge=0 still takes the legacy inv/pinv path and does not reach the Rust kernel. This matches the registry’s documented legacy-restoration contract.
• Concrete fix: None required.

Code Quality

No blocking findings. The Python fallback is factored into _ridge_solve_numpy() and the Rust symbol degrades cleanly when unavailable via diff_diff/_backend.py:L76-L96.

Performance

Finding: Remaining Prep Overhead Is Tracked

• Severity: P3 informational
• Location: TODO.md:L49, diff_diff/efficient_did_covariates.py:L1625-L1632
• Impact: The remaining zero_mask scan and omega_stack[rest] copy are performance-only overhead, not a correctness issue. This is already tracked in TODO.md.
• Concrete fix: None required for this PR. Future work can skip the fancy-index copy when rest.size == n.

Maintainability

No findings. The new Rust module is isolated in rust/src/batched_solve.rs, exported from rust/src/lib.rs:L57-L61, and covered by focused Python and Rust-side tests.

Tech Debt

No untracked blocking tech debt found. The only newly documented follow-up is the performance item in TODO.md:L49, which is acceptable under the deferred-work rules.

Security

No findings. I saw no unsafe Rust, secrets, credential material, or new external I/O paths in the changed files.

Documentation/Tests

No blocking findings. The PR adds targeted coverage for dispatch, Rust-vs-NumPy parity, bad-row recompute, singular pinv fallback, dtype dispatch refusal, split determinism, strided input, degenerate shapes, and shape validation at tests/test_efficient_did.py:L953-L1062 and tests/test_rust_backend.py:L3379-L3509.

I did not execute the test suite in this read-only review environment.