CelerisLab/tests/validation/time_window_dfg.py

"""Measurement-only A-versus-window timing for the matched shifted-MRT refined DFG setup.

Candidate source (window ledger) is recorded by checkout revision, dirty flag and file hashes;
the frozen refined-DFG manifests/results are NOT used or modified (they pin pre-window source).
Equal simulated duration, same fresh state and warm-up per variant, one observation interval =
80 root base steps (the runner's dense-phase interval):
  A   per base step: step(1) + read_raw_exchange, summed per interval (current runner workload)
  W   per interval: start_exchange_window(); step(80); read_exchange_window()
  W1  per interval: start_exchange_window(); 80 x step(1); read_exchange_window()
All variants keep observe_raw_exchange=True; the backend is built once with
observe_window_exchange=True (a closed window adds no launch). Populations (bitwise) and
per-interval total impulses are compared OUTSIDE timing; the impulse tolerance is a
conservative engineering tolerance (see INTERVAL_RTOL); the tighter substep-scaled engineering
tolerance is exercised by the capability test tests/validation/multigrid/test_exchange_window.py.
Wall time (perf_counter) with explicit stream synchronization is the only timing source.
"""

import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import statistics
import sys
import time

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import time_refined_dfg as t  # noqa: E402  (reuses run_refined_dfg setup via t.r)

r = t.r
INTERVAL = 80
VARIANTS = ("A", "W", "W1")
# Interval impulses differ only by FP64 summation order (per link over time vs per step over
# links). Engineering tolerance: 1e-9 of the interval's summed absolute per-step totals.
INTERVAL_RTOL = 1e-9
CANDIDATE_FILES = ("src/CelerisLab/lbm/refinement_stepper.py", "src/CelerisLab/lbm/raw_exchange.py",
                   "src/CelerisLab/lbm/refinement_tables.py", "src/CelerisLab/lbm/kernels/blocked.cu")


def components(obs):
    o = obs["total"]
    return np.array([o["impulse"][0], o["impulse"][1], o["angular_impulse"], o["mass_into_fluid"]])


def candidate_identity(source):
    from source_provenance import checkout_provenance
    return dict(checkout_provenance(source),
                sha256={f: hashlib.sha256((source / f).read_bytes()).hexdigest()
                        for f in CANDIDATE_FILES})


def run_variant(backend, variant, n_intervals, deadline):
    """Returns per-interval totals and, for A, per-interval summed |per-step totals|."""
    totals, scales = [], []
    for _ in range(n_intervals):
        if variant == "A":
            impulse, scale = np.zeros(4), np.zeros(4)
            for _ in range(INTERVAL):
                backend.step(1)
                c = components(backend.read_raw_exchange(0, level=1, window="last_base_step"))
                impulse += c
                scale += np.abs(c)
            totals.append(impulse)
            scales.append(scale)
        else:
            backend.start_exchange_window()
            if variant == "W":
                backend.step(INTERVAL)
            else:
                for _ in range(INTERVAL):
                    backend.step(1)
            totals.append(components(backend.read_exchange_window(0)))
        t.check_deadline(deadline)
    return np.array(totals), (np.array(scales) if scales else None)


def compare_intervals(reference, scales, totals):
    """Max violation ratio; <= 1 passes. Scale includes a floor so zero cancels cannot divide."""
    bound = INTERVAL_RTOL * scales + 1e-300
    return float(np.max(np.abs(totals - reference) / bound))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", type=Path, default=r.DEFAULT_SOURCE)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--base-steps", type=int, default=2000)
    p.add_argument("--warmup-base-steps", type=int, default=200)
    p.add_argument("--repeats", type=int, default=3)
    p.add_argument("--max-wall-seconds", type=float)
    p.add_argument("--plan", action="store_true")
    a = p.parse_args()
    if a.base_steps <= 0 or a.base_steps % INTERVAL or a.base_steps > 20000 or a.repeats < 1:
        p.error(f"base-steps must be a positive multiple of {INTERVAL} (<= 20000); repeats >= 1")
    a.source = a.source.resolve()
    identity = candidate_identity(a.source)
    plan = dict(collision="MRT", shifted_storage="as run_refined_dfg MRT config",
                candidate_source=identity, frozen_manifests_used=False, interval_base_steps=INTERVAL,
                base_steps=a.base_steps, warmup_base_steps=a.warmup_base_steps, repeats=a.repeats,
                variants=VARIANTS, tstar_window=a.base_steps * r.SPEED / 40.0,
                timer="time.perf_counter wall with explicit stream sync",
                excluded="construction/JIT, warm-up, health, downloads, parity/impulse comparison")
    if a.plan:
        print(json.dumps(plan, indent=2))
        return 0
    if a.max_wall_seconds is None or not math.isfinite(a.max_wall_seconds) or a.max_wall_seconds <= 0:
        p.error("GPU timing requires a finite positive --max-wall-seconds")
    if not os.environ.get("CUDA_VISIBLE_DEVICES") or "," in os.environ["CUDA_VISIBLE_DEVICES"]:
        p.error("Root must select exactly one verified idle GPU via CUDA_VISIBLE_DEVICES")
    if a.output.exists() and any(a.output.iterdir()):
        p.error("Output must be empty/new")
    a.output.mkdir(parents=True, exist_ok=True)
    deadline = time.monotonic() + a.max_wall_seconds
    h, _ = r.helpers(a.source)

    from contextlib import ExitStack
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    from CelerisLab.lbm.channel import channel_masks
    from CelerisLab.lbm.refinement import ACTIVE
    from CelerisLab.lbm.descriptors import SOLID

    rows, ref = [], {}
    n_intervals = a.base_steps // INTERVAL
    with ExitStack() as stack:
        cfg, pair, body, host = r.host_geometry(a.source, h, "MRT")
        context = stack.enter_context(CudaContext(0))
        backend = stack.enter_context(FixedPatchStepper(
            cfg, context=context, coarse_boundary="channel", refinements=[r.BOUNDS],
            level_blocks=[r.BASE_BLOCK, r.FINE_BLOCK], level_objects=[(), (body,)], observe_raw_exchange=True,
            observe_window_exchange=True))
        masks = [channel_masks(pair.coarse_shape)[0] & ~pair.covered,
                 (pair.fine_kind == ACTIVE) & ((host.flags.reshape(pair.fine_shape) & SOLID) == 0)]
        fresh = (r.initial_state(pair.coarse_shape, (0.0, 0.0), 1.0, root=True),
                 r.initial_state(pair.fine_shape, pair.fine_origin, 0.5))
        print(json.dumps(dict(status="CONSTRUCTED", case="refined")), flush=True)
        for repeat in range(a.repeats):
            for variant in VARIANTS:
                t.check_deadline(deadline)
                backend.initialize(*fresh)
                for _ in range(a.warmup_base_steps):
                    backend.step(1)
                warm = t.refined_health(backend, masks)
                if repeat == 0 and variant == "A":
                    print(json.dumps(dict(status="WARM_READY", case="refined")), flush=True)
                t.sync_refined(backend)
                start = time.perf_counter()
                totals, scales = run_variant(backend, variant, n_intervals, deadline)
                t.sync_refined(backend)
                seconds = time.perf_counter() - start
                end = t.refined_health(backend, masks)
                values = [np.array(v, copy=True) for v in backend._download_raw().values]
                if variant == "A" and repeat == 0:
                    ref.update(values=values, totals=totals, scales=scales)
                    parity, ratio = "identical", 0.0
                else:
                    parity = t.same_state(ref["values"], values) or "identical"
                    ratio = compare_intervals(ref["totals"], ref["scales"], totals)
                rows.append(dict(variant=variant, repeat=repeat, seconds=seconds,
                                 intervals=n_intervals, population_parity_vs_A=parity,
                                 interval_impulse_violation_ratio=ratio,
                                 health_after_warmup=warm, health_after_window=end))
                if parity != "identical" or not ratio <= 1:
                    raise RuntimeError(f"{variant} repeat {repeat}: parity {parity}, impulse ratio {ratio}")
    summary = {v: dict(median_seconds=statistics.median(x["seconds"] for x in rows if x["variant"] == v),
                       seconds=[x["seconds"] for x in rows if x["variant"] == v]) for v in VARIANTS}
    summary.update(W_over_A=summary["W"]["median_seconds"] / summary["A"]["median_seconds"],
                   W1_over_A=summary["W1"]["median_seconds"] / summary["A"]["median_seconds"],
                   interpretation="Measured comparative timings under GPU variability; W also "
                                  "batches step calls, W1 isolates the read pattern. Not a "
                                  "benchmark-throughput or causal claim.")
    (a.output / "window_timing.json").write_text(json.dumps(dict(plan=plan, rows=rows, summary=summary),
                                                           indent=2))
    print(json.dumps(summary, indent=2), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())