Skip to content

perf(efficient_did): cross-cell kcov-table hoisting + per-group W lifecycle (100k: 86 -> 18 min)#629

Merged
igerber merged 1 commit into
mainfrom
perf/efficient-did-kcov-hoist
Jul 7, 2026
Merged

perf(efficient_did): cross-cell kcov-table hoisting + per-group W lifecycle (100k: 86 -> 18 min)#629
igerber merged 1 commit into
mainfrom
perf/efficient-did-kcov-hoist

Conversation

@igerber

@igerber igerber commented Jul 6, 2026

Copy link
Copy Markdown
Owner

Summary

  • Hoists the conditional-path kernel-covariance tables across (g, t) cells: every Omega* term is s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | group) keyed only by wide-outcome columns, so the tables the v3.7 tiled pass rebuilt per cell dedup to one table of distinct product columns per comparison group per unit-tile (~26x fewer kernel GEMM columns on a PT-All fit).
  • Difference keys are ORDERED (minuend, subtrahend) and never sort-canonicalized (sorting would sign-flip covariances for cells whose pre-periods straddle t under PT-All); only the symmetric product-pair swap is canonicalized.
  • Kernel weight matrices are built and freed one group at a time, so the tile memory budget is governed by the largest single group instead of the sum (proportionally fatter tiles at large n).
  • Each cell's Omega* is gathered from the group tables as a triu strip through compact per-cell column slices in the legacy per-entry operation order, then mirrored — value-exact vs the legacy loop (locked by assert_array_equal); masked entries resolve to an all-zero column.
  • Measured (3-rep medians, 20 periods / 5 cohorts / 5 covariates, aggregate="all"): 2k 7.7 -> 7.0s, 10k 57.2 -> 35.7s (1.6x), 30k 358 -> 129s (2.8x), 100k 85.6 min -> 17.8 min (~4.8x) at the same ~4 GB peak RSS; the kernel-covariance stage itself is 40x faster (21.2 -> 0.5s at 10k). Results move only at floating-point reassociation level (post-treatment cells ~1e-12 relative, overall ATT ~1e-13); the no-covariates path is byte-identical; omega_ridge=0 still routes the entire legacy path; a pt_assumption="post" covariate fit builds no tables at all.
  • Follow-up TODO row filed: post-hoist the top stage is _ridge_solve_weights (batched LU; Rust batched-Cholesky candidate).

Methodology references (required if estimator / math changes)

  • Method name(s): EfficientDiD (efficient staggered DiD), conditional Omega* construction (Eq 3.12) feeding the efficient weights (Eq 4.3)
  • Paper / source link(s): Chen, X., Sant'Anna, P. H. C., & Xie, H. (2025). Efficient Difference-in-Differences and Event Study Estimators. arXiv:2506.17729
  • Any intentional deviations from the source (and why): None — pure implementation change (same five-term construction, same values up to floating-point reassociation). The cross-cell hoisting is documented in the Omega* ridge Note in docs/methodology/REGISTRY.md.

Validation

  • Tests added/updated: tests/test_methodology_efficient_did.py (new TestKcovTableHoisting: per-group table oracle vs the per-cell _kcov_batch construction, unweighted + survey-weighted, rel 1e-12; reversed-orientation ordered-key lock; multi-y1_col defensive case; value-exact assembly-vs-legacy-loop; H==1 no-keys), tests/test_efficient_did.py (build-count-proportional-to-tiles proof, pt_assumption="post" zero-builds guard, mid-size tile twin with forced product chunking). All existing locks (dense-reference oracle, 1-byte tile twins, 1-ulp stability, HRS anchors) pass unmodified. Full suite: 8507 passed (1 known pre-existing order-dependent dCDH failure); targeted suites 461 passed on both backends; tutorial 15 re-executed clean via nbmake.
  • Backtest / simulation / notebook evidence (if applicable): perf arms in docs/performance-plan.md (BEFORE = pristine main fcd7768 via baseline-worktree PYTHONPATH; 3-rep medians; ATT/SE delta report included).

Security / privacy

  • Confirm no secrets/PII in this PR: Yes

Generated with Claude Code

🤖 Generated with Claude Code

https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr

…ecycle (100k: 86->18 min)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr
@github-actions

github-actions Bot commented Jul 6, 2026

Copy link
Copy Markdown

Overall Assessment

✅ Looks good. No unmitigated P0 or P1 findings.

Executive Summary

  • Affected method: EfficientDiD covariate path, conditional Omega* construction feeding efficient weights.
  • The hoisting is documented in docs/methodology/REGISTRY.md:L1246-L1258 under a **Note:**, and appears to preserve the five-term Eq. 3.12 construction rather than changing identification.
  • No new inline inference, NaN/SE, default-behavior, or parameter-propagation anti-patterns were introduced.
  • Tests cover the key contracts: ordered difference keys, survey-weighted tables, H==1 no-table path, chunking, tiling, and dense-reference equivalence.
  • I could not execute tests locally because pytest is not installed in this runner.

Methodology

No findings.

Severity: N/A
Impact: The affected code in diff_diff/efficient_did_covariates.py:L1143-L1591 rewrites how kernel-covariance terms are shared and gathered, but it keeps the same term scaling and group semantics as the legacy conditional Omega* path. The documented registry note explicitly classifies cross-cell table hoisting as an implementation change, not a methodology deviation.
Concrete fix: None.

Code Quality

No findings.

Severity: N/A
Impact: The helper split is readable and locally documented. The assembly logic preserves ordered difference keys and uses explicit zero-column masking in diff_diff/efficient_did_covariates.py:L1216-L1300 and L1323-L1338.
Concrete fix: None.

Performance

P3-informational, tracked in TODO.md.

Impact: The PR resolves the prior kernel-covariance bottleneck and adds a new tracked follow-up for _ridge_solve_weights batched LU/possible Cholesky work in TODO.md:L47-L50. This is performance-only and properly tracked.
Concrete fix: No action required for this PR.

Maintainability

No findings.

Severity: N/A
Impact: The new registry/table helpers are private, scoped to the fused conditional path, and have dedicated methodology tests in tests/test_methodology_efficient_did.py:L476-L670.
Concrete fix: None.

Tech Debt

No blocking findings.

Severity: P3-informational
Impact: The remaining Rust batched-Cholesky optimization is explicitly tracked in TODO.md:L47-L50, so it is accepted deferred work.
Concrete fix: No action required.

Security

No findings.

Severity: N/A
Impact: The diff is numerical code, tests, changelog, TODO, and methodology/performance docs. I did not see secrets or security-sensitive behavior introduced.
Concrete fix: None.

Documentation/Tests

P3-informational.

Impact: Coverage is strong for the new hoisting contracts, especially tests/test_efficient_did.py:L797-L882 and tests/test_methodology_efficient_did.py:L556-L670. One optional gap remains: there is no direct control_group="last_cohort" plus covariates test exercising the hoisted conditional path, though the mask logic appears unchanged and correct.
Concrete fix: Optional: add a small overidentified covariate test with no never-treated units, control_group="last_cohort", and pt_assumption="all", asserting finite estimates and agreement with the dense reference.

Verification note: targeted pytest command could not run because /bin/bash: pytest: command not found.

@igerber igerber added the ready-for-ci Triggers CI test workflows label Jul 6, 2026
@igerber igerber merged commit 81d53f5 into main Jul 7, 2026
35 of 36 checks passed
@igerber igerber deleted the perf/efficient-did-kcov-hoist branch July 7, 2026 00:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready-for-ci Triggers CI test workflows

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant