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

CHANGELOG.md

@@ -127,6 +127,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   supersedes it.
 
 ### Changed
+- **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,
+  10-cohort, 100k-unit panel (390 cells, 999 draws) that was ~95% of bootstrap wall time as
+  pure memory traffic. All perturbation columns (per-cell IFs, overall combined IF,
+  per-event-time combined IFs) are now scattered into a column-tiled influence matrix
+  (byte-capped tiles; each tile replays the bit-identical weight stream via RNG state
+  snapshot/restore) and consumed by one BLAS GEMM per weight block. Measured (5-rep medians,
+  rust backend + Accelerate): bootstrap fits 24.5s → 4.3s (5.7x) at the 4M-row flagship,
+  5.6s → 0.31s (18.3x) no-covariate and 6.5s → 1.2s (5.5x) dr at 2M rows, 4.2s → 1.5s (2.7x)
+  survey-PSU; peak memory flat or lower. Point estimates are bit-identical (the bootstrap
+  never feeds them); bootstrap SEs/CIs/p-values match the previous loop to BLAS
+  reassociation (≤ ~1e-15 relative — far below bootstrap Monte-Carlo error), the same
+  numerics posture as the existing draw-axis chunking. EfficientDiD's unweighted path folds
+  its `1/n` prefactor into the influence column (`W @ (eif/n)` instead of `(W @ eif)/n`) —
+  the same reassociation-level equivalence.
 - **CallawaySantAnna aggregation SE assembly rewritten to O(n_units)** (also inherited by
   `StaggeredTripleDifference` aggregation). The combined influence-function assembly behind
   simple/event-study/group aggregated SEs — previously ~56-85% of analytical fit time at

TODO.md

@@ -47,6 +47,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
+| `EfficientDiD` analytical stage has a pathological scaling hot loop: 339s for ONE fit at just 2k units × 20 periods with 5 covariates (15+ CPU-min at 10k units; memory pressure at 100k), while CS handles 100k units in ~2s. Process sampling attributes it to a broadcasted NumPy subtract/multiply inside a Python-level loop (EIF/weights construction). Profile properly and vectorize; until then EfficientDiD is impractical beyond a few thousand units. Found while benchmarking the Phase 2 bootstrap rewrite (its bootstrap stage is 0.018s at 2k units — the analytical stage is the whole problem). | `efficient_did.py` / `efficient_did_weights.py` | CS-scaling | Heavy | Medium |
 | `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/bootstrap_chunking.py

@@ -23,11 +23,27 @@
 - The **NumPy** fallback draws the matrix row-major from the ``Generator``
   stream, so consuming it in contiguous, in-order blocks from the same generator
   reproduces the identical sequence.
+
+The *influence* side can be tiled too. When a fit needs many perturbation
+columns (per-cell IFs, the overall combined IF, per-event-time combined IFs),
+computing them one column at a time degenerates into per-column slicing/GEMV
+over the weight blocks -- memory-bandwidth-bound, not FLOP-bound. Instead,
+:func:`tiled_if_matmul` scatters the columns into an ``(n_units, tile_cols)``
+buffer and runs one BLAS GEMM per (weight block, column tile). The column axis
+is tiled under a byte cap, and each column tile makes its own full pass over
+the weight stream via :class:`ReplayableWeightStream`, which snapshots the
+generator state at construction and restores it per pass -- so every pass sees
+the bit-identical weight stream on both backends (Rust re-draws the same
+``base_seed`` from the restored state; the NumPy fallback replays the stream
+itself; the survey stratified branch regenerates its full PSU matrix from the
+restored state). GEMM results match the per-column GEMV path to within BLAS
+reassociation, same as the draw-axis chunking above.
 """
 
 from __future__ import annotations
 
-from typing import Iterator, Optional, Tuple
+from collections import abc
+from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -183,3 +199,134 @@ def _generated() -> Iterator[Tuple[int, np.ndarray]]:
             yield chunk_start, block
 
     return psu_ids, _generated()
+
+
+# Byte ceiling for the ``(n_units, tile_cols)`` float64 scatter buffer used by
+# ``tiled_if_matmul``. Sibling of ``_TARGET_BLOCK_BYTES``; together they bound
+# the kernel's live intermediates at ~384 MB regardless of n_units or the
+# number of perturbation columns. Read at CALL time inside ``tiled_if_matmul``
+# (never bound as a def-time default) so tests can monkeypatch it.
+_TARGET_TILE_BYTES = 128 * 1024 * 1024
+
+# One perturbation column of ``tiled_if_matmul``'s influence matrix: a list of
+# (index, values) contributions. ``index=None`` assigns the full-width dense
+# ``values``; otherwise ``values`` is scattered at ``index``. A column may also
+# be a zero-argument callable returning such a list, materialized only when its
+# tile is filled (keeps peak memory at one O(n_units) temporary for callers
+# whose columns are computed on the fly).
+IFColumn = List[Tuple[Optional[np.ndarray], np.ndarray]]
+IFColumnSpec = Union[IFColumn, Callable[[], IFColumn]]
+
+
+class ReplayableWeightStream:
+    """Multiplier-weight block stream that can be re-iterated bit-identically.
+
+    Wraps a weight-block iterator factory together with the ``Generator`` that
+    feeds it, snapshotting the generator state at construction. Every
+    ``__iter__`` restores the snapshot before invoking the factory, so each
+    pass yields the exact same ``(chunk_start, block)`` sequence:
+
+    - **Rust backend**: ``iter_weight_blocks`` draws its per-run ``base_seed``
+      from the rng lazily at the first ``next()``; the restored state
+      reproduces the same base seed, and rows are seeded absolutely from it.
+    - **NumPy fallback**: generation consumes the rng stream directly; the
+      restored state replays the identical stream.
+    - **Survey stratified branch**: ``iter_survey_multiplier_weight_blocks``
+      generates its full PSU matrix eagerly at factory-call time -- from the
+      restored state, so per-pass recreation is also bit-identical.
+
+    The snapshot MUST precede any rng-consuming call: construct this object
+    before probing generators that draw at call time (the survey stratified
+    branch draws eagerly; resolve ``psu_ids`` from the design's ``psu`` array
+    directly instead of via a probe call). After the final pass the rng ends
+    exactly where a single pass would leave it, so downstream consumers of the
+    same generator are unaffected by the number of passes.
+    """
+
+    def __init__(
+        self,
+        make_iter: Callable[[np.random.Generator], Iterator[Tuple[int, np.ndarray]]],
+        rng: np.random.Generator,
+    ) -> None:
+        self._make_iter = make_iter
+        self._rng = rng
+        self._state = rng.bit_generator.state
+
+    def __iter__(self) -> Iterator[Tuple[int, np.ndarray]]:
+        self._rng.bit_generator.state = self._state
+        return self._make_iter(self._rng)
+
+
+def tiled_if_matmul(
+    weight_stream: Iterable[Tuple[int, np.ndarray]],
+    n_bootstrap: int,
+    n_units: int,
+    columns: Sequence[IFColumnSpec],
+    *,
+    tile_bytes: Optional[int] = None,
+) -> np.ndarray:
+    """Compute ``W @ M`` for the multiplier bootstrap without materializing
+    ``W`` (streamed in row blocks) or ``M`` (scattered in column tiles).
+
+    ``W`` is the ``(n_bootstrap, n_units)`` weight matrix yielded by
+    ``weight_stream`` as ``(chunk_start, block)`` row blocks; ``M`` is the
+    ``(n_units, len(columns))`` influence matrix described by ``columns``
+    (see :data:`IFColumnSpec`). Returns the dense ``(n_bootstrap, n_cols)``
+    product. Columns are tiled under ``tile_bytes`` (default: module attribute
+    ``_TARGET_TILE_BYTES``, resolved at call time); each tile is scattered
+    once into an F-order buffer and then consumed by one BLAS GEMM per weight
+    block. With more than one tile, ``weight_stream`` is iterated once per
+    tile and therefore must be re-iterable with an identical stream --
+    :class:`ReplayableWeightStream` provides exactly that.
+
+    Scatter contract: contributions are written by plain assignment, so index
+    arrays must be unique within each contribution and the contributions of a
+    single column must be mutually disjoint (duplicate or overlapping indices
+    keep the last write instead of summing like ``W[:, idx] @ values`` would).
+    All in-repo callers satisfy this by construction (treated/control unit
+    sets are disjoint; dense columns are single-contribution).
+
+    A non-finite influence value poisons only its own output column (every
+    output element is an independent dot product), matching the per-column
+    GEMV behavior this kernel replaces. Results match the per-column GEMV
+    path to within BLAS reassociation (~1 ULP), not bit-for-bit.
+    """
+    n_cols = len(columns)
+    out = np.empty((n_bootstrap, n_cols))
+    if n_cols == 0 or n_bootstrap == 0:
+        return out
+    if tile_bytes is None:
+        tile_bytes = _TARGET_TILE_BYTES
+    tile_cols = int(max(1, min(n_cols, tile_bytes // (max(1, n_units) * 8))))
+
+    # A single-pass iterator silently exhausts after the first tile, leaving
+    # uninitialized rows in `out` -- fail loudly instead. isinstance (not an
+    # iter() probe) so re-iterable streams see no extra pass.
+    if tile_cols < n_cols and isinstance(weight_stream, abc.Iterator):
+        raise ValueError(
+            "tiled_if_matmul needs more than one column tile but weight_stream "
+            "is a single-pass iterator; pass a re-iterable stream "
+            "(ReplayableWeightStream) so each tile can replay the weights."
+        )
+
+    buf = np.zeros((n_units, tile_cols), order="F")
+    for lo in range(0, n_cols, tile_cols):
+        hi = min(lo + tile_cols, n_cols)
+        width = hi - lo
+        if lo > 0:
+            buf[:, :width] = 0.0
+        for c in range(lo, hi):
+            spec = columns[c]
+            if callable(spec):
+                spec = spec()
+            col = buf[:, c - lo]
+            for idx, values in spec:
+                if idx is None:
+                    np.copyto(col, values)
+                else:
+                    col[idx] = values
+        tile = buf[:, :width]
+        with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
+            for chunk_start, w_block in weight_stream:
+                out[chunk_start : chunk_start + w_block.shape[0], lo:hi] = w_block @ tile
+    return out

diff_diff/efficient_did_bootstrap.py

@@ -13,9 +13,11 @@
 import numpy as np
 
 from diff_diff.bootstrap_chunking import (
+    ReplayableWeightStream,
     compute_block_size,
     iter_survey_multiplier_weight_blocks,
     iter_weight_blocks,
+    tiled_if_matmul,
 )
 from diff_diff.bootstrap_utils import (
     compute_effect_bootstrap_stats as _compute_effect_bootstrap_stats_func,
@@ -97,7 +99,6 @@ def _run_multiplier_bootstrap(
         rng = np.random.default_rng(self.seed)
 
         gt_pairs = list(group_time_effects.keys())
-        n_gt = len(gt_pairs)
 
         # Original ATTs (independent of the draws; referenced per block below).
         original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs])
@@ -120,13 +121,15 @@ def _run_multiplier_bootstrap(
             # at a time (unstratified designs tile the generation; stratified
             # designs have few PSUs and fall back to full generation + slicing).
             _block_size = compute_block_size(n_units, self.n_bootstrap)
-            psu_ids, _psu_blocks = iter_survey_multiplier_weight_blocks(
-                self.n_bootstrap,
-                resolved_survey,
-                self.bootstrap_weights,
-                rng,
-                block_size=_block_size,
-            )
+            # Resolve psu_ids WITHOUT calling the generator: the stratified
+            # branch draws from the rng eagerly at call time, and the
+            # replayable stream below must snapshot the rng state before any
+            # draw. Duplicates the rng-free resolution both generator branches
+            # use (np.unique / np.arange).
+            if resolved_survey.psu is not None:
+                psu_ids = np.unique(resolved_survey.psu)
+            else:
+                psu_ids = np.arange(len(resolved_survey.weights))
             # Single-cluster (G<2) survey-PSU multiplier bootstrap collapses
             # to constant multiplier draws → BLAS roundoff produces ≈0
             # variance (NOT NaN). Downstream zero-SE guards check exact 0 and
@@ -163,29 +166,51 @@ def _run_multiplier_bootstrap(
                 np.array_equal(unit_to_psu_col, np.arange(n_units))
             )
 
-            def _weight_blocks() -> Iterator[Tuple[int, np.ndarray]]:
-                for _cs, _psu_block in _psu_blocks:
-                    if _psu_is_identity:
-                        yield _cs, _psu_block
-                    else:
-                        yield _cs, _psu_block[:, unit_to_psu_col]
+            # Factory recreating the PSU generation + unit expansion per pass.
+            def _make_weight_iter(
+                rng_: np.random.Generator,
+            ) -> Iterator[Tuple[int, np.ndarray]]:
+                _, _psu_blocks = iter_survey_multiplier_weight_blocks(
+                    self.n_bootstrap,
+                    resolved_survey,
+                    self.bootstrap_weights,
+                    rng_,
+                    block_size=_block_size,
+                )
+
+                def _expanded() -> Iterator[Tuple[int, np.ndarray]]:
+                    for _cs, _psu_block in _psu_blocks:
+                        if _psu_is_identity:
+                            yield _cs, _psu_block
+                        else:
+                            yield _cs, _psu_block[:, unit_to_psu_col]
+
+                return _expanded()
 
-            weight_blocks: Iterator[Tuple[int, np.ndarray]] = _weight_blocks()
         elif cluster_indices is not None and n_clusters is not None:
             # Cluster-level weights, expanded to unit level per block via the
             # helper's expand_index (block[:, cluster_indices]).
-            weight_blocks = iter_weight_blocks(
-                self.n_bootstrap,
-                n_clusters,
-                self.bootstrap_weights,
-                rng,
-                expand_index=cluster_indices,
-            )
+            def _make_weight_iter(
+                rng_: np.random.Generator,
+            ) -> Iterator[Tuple[int, np.ndarray]]:
+                return iter_weight_blocks(
+                    self.n_bootstrap,
+                    n_clusters,
+                    self.bootstrap_weights,
+                    rng_,
+                    expand_index=cluster_indices,
+                )
+
         else:
             # Standard unit-level weights, generated one row-block at a time.
-            weight_blocks = iter_weight_blocks(
-                self.n_bootstrap, n_units, self.bootstrap_weights, rng
-            )
+            def _make_weight_iter(
+                rng_: np.random.Generator,
+            ) -> Iterator[Tuple[int, np.ndarray]]:
+                return iter_weight_blocks(self.n_bootstrap, n_units, self.bootstrap_weights, rng_)
+
+        # Re-iterable stream: each column tile of the fused perturbation GEMM
+        # below makes its own full pass over the bit-identical weight stream.
+        weight_stream = ReplayableWeightStream(_make_weight_iter, rng)
 
         # eif SCALING is a SEPARATE axis from the weight PATH: it is keyed on
         # unit_level_weights (set whenever a SurveyDesign was passed — including a
@@ -196,24 +221,32 @@ def _weight_blocks() -> Iterator[Tuple[int, np.ndarray]]:
         _has_unit_weights = unit_level_weights is not None
         _total_w = float(np.sum(unit_level_weights)) if _has_unit_weights else 1.0
 
-        # Pre-allocate the small (n_bootstrap, n_gt) output; only this persists.
-        # Each weight block fills its rows of every column, then is discarded —
-        # peak memory is capped at O(block x n_units). The per-(g,t) scaled EIF is
-        # recomputed per block as a single O(n_units) temporary (not cached across
-        # all n_gt cells), so the perturbation adds no O(n_gt x n_units) allocation
-        # that would erode the memory win on weighted panels. The aggregations
-        # below (overall, event study, group) re-aggregate these columns and never
-        # touch the weight matrix.
-        bootstrap_atts = np.empty((self.n_bootstrap, n_gt))
-        for _chunk_start, _w_block in weight_blocks:
-            _rows = slice(_chunk_start, _chunk_start + _w_block.shape[0])
-            for j, gt in enumerate(gt_pairs):
-                with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-                    if _has_unit_weights:
-                        perturbation = _w_block @ (unit_level_weights * eif_by_gt[gt] / _total_w)
-                    else:
-                        perturbation = (_w_block @ eif_by_gt[gt]) / n_units
-                bootstrap_atts[_rows, j] = original_atts[j] + perturbation
+        # Fused perturbation GEMM over the replayable weight stream, one column
+        # per (g,t). Columns are LAZY callables materializing each scaled EIF
+        # only when its tile is filled — one O(n_units) temporary at a time, so
+        # the perturbation still adds no O(n_gt x n_units) allocation that
+        # would erode the memory win on weighted panels (tiles are capped by
+        # _TARGET_TILE_BYTES). The unweighted path folds its 1/n prefactor into
+        # the column (W @ (eif/n) instead of (W @ eif)/n) — a pure BLAS
+        # reassociation-level change. The aggregations below (overall, event
+        # study, group) re-aggregate these columns and never touch the weight
+        # matrix.
+        def _scaled_eif_column(gt: Tuple[Any, Any]):
+            def _make() -> List[Tuple[Optional[np.ndarray], np.ndarray]]:
+                if _has_unit_weights:
+                    return [(None, unit_level_weights * eif_by_gt[gt] / _total_w)]
+                return [(None, eif_by_gt[gt] / n_units)]
+
+            return _make
+
+        perturbations = tiled_if_matmul(
+            weight_stream,
+            self.n_bootstrap,
+            n_units,
+            [_scaled_eif_column(gt) for gt in gt_pairs],
+        )
+        with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
+            bootstrap_atts = original_atts[None, :] + perturbations
 
         # Post-treatment mask — also exclude NaN effects
         post_mask = np.array(