perf(bootstrap): fused column-tiled scatter-GEMM multiplier bootstrap (CS + EfficientDiD)

[igerber]

Summary

• Replace the CallawaySantAnna multiplier bootstrap's per-cell weight-matrix slicing loop (two fancy-index copies + two GEMVs per (g,t) cell per weight block - measured 95% of bootstrap wall time as pure memory traffic) with one column-tiled BLAS GEMM over a fused influence matrix [per-cell IFs | overall combined IF | event-study combined IFs].
• New shared kernel in bootstrap_chunking.py: tiled_if_matmul (column tiles under a call-time byte cap; assignment-scatter contract documented; loud ValueError if a single-pass iterator meets multiple tiles) + ReplayableWeightStream (weight-block stream re-iterable bit-identically via RNG state snapshot/restore - rust per-row seeding, numpy stream, and the stratified-survey eager-generation branch all replay exactly).
• EfficientDiD's per-cell full-width GEMV bootstrap routes through the same kernel via lazily materialized scaled-EIF columns (preserves its documented avoidance of the O(n_gt x n_units) allocation). HAD already consumed whole-matrix GEMMs and is untouched.
• psu_ids + the G<2 guards now resolve rng-free (np.unique/np.arange) outside the replay factories; CS group aggregation moves post-loop (reads only completed perturbed cell draws).
• Measured (5-rep medians, rust backend + Accelerate; BEFORE = pristine origin/main worktree): bootstrap fits 24.5s -> 4.3s (5.7x) at the 4M-row 390-cell flagship; 5.6s -> 0.31s (18.3x) no-cov and 6.5s -> 1.2s (5.5x) dr at 2M rows; 4.2s -> 1.5s (2.7x) survey-PSU (non-identity map); RCS 1.1x (bootstrap is a small share there). Peak RSS flat or lower. Full tables in docs/performance-plan.md.
• ATT point estimates bit-identical on every scenario (the bootstrap never feeds them); bootstrap SEs/CIs/p-values match the previous loop to BLAS reassociation (3e-16..1.2e-15 relative - far below Monte-Carlo error), the same numerics posture as the existing draw-axis chunking. EfficientDiD's unweighted path folds its 1/n prefactor into the column (W @ (eif/n) vs (W @ eif)/n) - same reassociation-level equivalence, called out in comments/CHANGELOG.
• New TODO row (found while benchmarking, out of scope here): EfficientDiD's ANALYTICAL stage has a pre-existing pathological hot loop (339s at 2k units x 20p with 5 covariates; sampled to a broadcasted subtract/multiply in a Python-level loop).

Methodology references (required if estimator / math changes)

• Method name(s): Callaway & Sant'Anna (2021) multiplier bootstrap; EfficientDiD multiplier bootstrap. Numerical kernel refactor only - the perturbation algebra, weight generation, and aggregation are unchanged.
• Paper / source link(s): Callaway & Sant'Anna (2021), Journal of Econometrics 225(2); REGISTRY.md "Bootstrap, Survey Bootstrap" sections.
• Any intentional deviations from the source (and why): None. Results match the previous implementation to BLAS reassociation (~1e-15 relative); point estimates bit-identical.

Validation

• Tests added/updated: tests/test_bootstrap_chunking.py (kernel oracle vs the old per-column GEMV semantics; replay determinism across rademacher/mammen/webb; counting guard proving multi-tile actually replays; single-pass-iterator guard; tile-forced end-to-end invariance twins for CS unit/cluster/survey-PSU/STRATIFIED-survey and all four EfficientDiD weight paths), tests/test_staggered.py (NaN-cell bootstrap isolation). Targeted sweep 951 passed on BOTH backends (rust + DIFF_DIFF_BACKEND=python); full suite 8405 passed with one known pre-existing order-dependent dCDH failure (passes standalone; tracked in TODO.md since fix(staggered): CS panel reg/ipw influence functions + per-cell SEs to DRDID parity #619). No existing test expectation changed or weakened.
• Backtest / simulation / notebook evidence (if applicable): A/B perf arms vs pristine origin/main recorded in docs/performance-plan.md (Phase 2 section), including the stage-level split that motivated the rewrite and a measured-and-rejected scipy-sparse variant.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected methods: CallawaySantAnna multiplier bootstrap, CallawaySantAnna survey/PSU bootstrap, and EfficientDiD multiplier bootstrap.
• Methodology matches the registry: this is a GEMV-to-GEMM numerical refactor, not a change to IF algebra, weighting, assumptions, or inference semantics.
• NaN/Inf handling remains centralized through bootstrap stat helpers; no new inline effect / se anti-pattern found.
• Survey PSU ordering appears preserved (np.unique / np.arange matches the generator helper).
• Two P3 items noted: one RCS tile-invariance test gap, one minor doc-deps wording issue.

Methodology

No blocking findings.

The PR preserves the documented CS contract that bootstrap perturbs per-cell IFs and combined IF/WIF aggregations, and preserves the EfficientDiD documented multiplier-on-EIF implementation choice. The new tiled_if_matmul path in diff_diff/bootstrap_chunking.py:L260-L331, CS wiring in diff_diff/staggered_bootstrap.py:L468-L518, and EfficientDiD wiring in diff_diff/efficient_did_bootstrap.py:L224-L249 are equivalent numerical reductions, not methodology deviations.

Code Quality

No findings.

The replayable stream design is explicit and guarded for single-pass iterators in diff_diff/bootstrap_chunking.py:L221-L257 and L302-L310. The assignment-scatter contract is documented and satisfied by current callers.

Performance

No blocking findings.

The performance-oriented rewrite is consistent with the stated goal. The new EfficientDiD analytical hot-loop concern is not introduced by this PR’s bootstrap rewrite.

Maintainability

No findings.

The shared kernel reduces duplicate bootstrap matmul logic across CS and EfficientDiD.

Tech Debt

• Severity: P3
  Impact: The PR adds a tracked EfficientDiD analytical-stage scaling issue in TODO.md:L50. This is performance debt, not a silent correctness issue.
  Concrete fix: Profile efficient_did.py / efficient_did_weights.py EIF construction and vectorize the identified broadcasted Python-loop path when prioritized.

Security

No findings.

No suspicious secret patterns were found in the changed files.

Documentation/Tests

• Severity: P3
  Impact: CS repeated cross-section bootstrap is affected by the new tiled kernel, but the new tile-forced end-to-end tests cover panel/unit, cluster, survey PSU, stratified survey, and EfficientDiD paths, not RCS. Existing RCS bootstrap tests exercise default behavior in tests/test_staggered_rc.py:L286-L301 and L416-L443, but not forced multi-tile replay.
  Concrete fix: Add a small CallawaySantAnna(panel=False, n_bootstrap=..., cband=False, aggregate="all") tile-invariance test with _TARGET_TILE_BYTES=1, comparing default vs forced-tile SE/CI outputs.

• Severity: P3
  Impact: docs/doc-deps.yaml:L946-L952 reads as if the fused tiled_if_matmul kernel is shared with HAD, while HAD appears untouched and only shares the broader chunking machinery.
  Concrete fix: Reword to distinguish “draw-block chunking shared with HAD” from “fused column-tiled kernel used by CS and EfficientDiD.”

Validation note: git diff --check passed. I attempted a tiny in-memory RCS smoke check, but the review environment is missing numpy, so runtime tests could not be executed here.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 6070479dce6bbe31d3de8d0ae72fd19b06494887

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected methods: CallawaySantAnna multiplier bootstrap, including survey/PSU and repeated-cross-section paths, and EfficientDiD multiplier bootstrap.
• Methodology appears preserved: the PR changes the numerical contraction from per-column GEMV to tiled GEMM, while retaining the same IF/EIF inputs, weights, aggregation rules, and inference helpers.
• Prior P3 items are addressed: RCS forced-tile coverage was added, and docs/doc-deps.yaml now clearly says HAD does not use the new tiled kernel.
• No new inline inference or partial NaN guard anti-pattern was introduced in the changed files.
• One P3 tracked performance-debt item remains: the pre-existing EfficientDiD analytical hot loop is now recorded in TODO.md.

Methodology

No blocking findings.

Affected source contracts checked: CS bootstrap/WIF contract in docs/methodology/REGISTRY.md:L536-L542, CS non-finite bootstrap contract in docs/methodology/REGISTRY.md:L604-L609, EfficientDiD bootstrap/EIF contract in docs/methodology/REGISTRY.md:L1150-L1159 and docs/methodology/REGISTRY.md:L1185-L1186.

• Severity: P3-informational
  Impact: EfficientDiD continues to use multiplier bootstrap on stored EIF values rather than the nonparametric clustered bootstrap mentioned in the paper. This is already documented in the registry as an implementation note and is not a defect.
  Concrete fix: No action required.

The new CS path builds perturbation columns for per-cell IFs, overall combined IF, and event-study combined IFs in diff_diff/staggered_bootstrap.py:L468-L493, then reconstructs bootstrap draws and delegates inference to existing helpers in diff_diff/staggered_bootstrap.py:L495-L620. EfficientDiD likewise preserves the old scaling, including survey-weighted w_i * eif_i / sum(w) and unweighted eif_i / n, in diff_diff/efficient_did_bootstrap.py:L224-L249.

Code Quality

No findings.

The reusable kernel documents the assignment-scatter contract and guards multi-tile use with single-pass iterators in diff_diff/bootstrap_chunking.py:L282-L310. Current in-repo callers satisfy the disjoint-index contract.

Performance

No blocking findings.

The change is performance-oriented and avoids the previous per-cell weight slicing loop. The remaining EfficientDiD analytical scaling issue is not introduced by this PR.

Maintainability

No findings.

The shared tiled_if_matmul / ReplayableWeightStream path reduces duplicated bootstrap matmul logic across CS and EfficientDiD.

Tech Debt

• Severity: P3
  Impact: The PR adds a tracked EfficientDiD analytical-stage performance issue in TODO.md:L50. This is performance debt, not silent statistical incorrectness.
  Concrete fix: Profile efficient_did.py / efficient_did_weights.py EIF construction and vectorize the identified broadcasted loop when prioritized.

Security

No findings.

No suspicious secret patterns were found in the changed files.

Documentation/Tests

No blocking findings.

Prior review items are resolved: RCS forced-tile coverage is now present in tests/test_bootstrap_chunking.py:L569-L608, and the doc-deps note now distinguishes shared draw-block chunking from the CS/EfficientDiD-only tiled kernel in docs/doc-deps.yaml:L946-L952.

Validation note: git diff --check passed. I did not run the pytest suite in this read-only review environment.