CelerisLab/tests/validation/time_refined_dfg.py

"""Measurement-only wall timing for the matched MRT refined DFG case. No numerical edits.

Variants over EQUAL simulated duration (one refined base step = two uniform D80 steps):
  A  refined: step(1) + read_raw_exchange per base step (as run_refined_dfg)
  B  refined: step(1) only                       -- timing-only, no force history
  C  refined: one step(N) call                   -- timing-only, no force history
  E  uniform: sim.run(<=100) + read_force (as run_refined_dfg)
  E' uniform: one sim.run(2N) call                -- timing-only, no force history
A-B and B-C are measured comparative differences, not exact causal decompositions.
Each repeat starts from the same fresh initial state with the same warm-up; construction,
compilation and parity checks are outside timed regions. FixedPatchStepper.step() synchronizes
its stream after every base interval and Simulation.run() synchronizes before returning, so the
timer starts and stops with no GPU work in flight. Wall time (perf_counter) is the only source.
"""

import argparse
import json
import math
import os
from pathlib import Path
import statistics
import sys
import time
from tempfile import TemporaryDirectory

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import run_refined_dfg as r  # noqa: E402  (published runner: same setup functions)

VARIANTS = {
    "A": dict(case="refined", force_history_valid=True),
    "B": dict(case="refined", force_history_valid=False, timing_only=True),
    "C": dict(case="refined", force_history_valid=False, timing_only=True),
    "E": dict(case="uniform", force_history_valid=True),
    "E'": dict(case="uniform", force_history_valid=False, timing_only=True),
}


def schedule(n_base, warmup_base):
    """Steps per case for equal simulated duration (refined dt_root = 2 dt_uniform)."""
    if n_base <= 0 or warmup_base < 0:
        raise ValueError("n_base must be positive and warmup_base nonnegative")
    return dict(refined=(warmup_base, n_base), uniform=(2 * warmup_base, 2 * n_base),
                tstar_window=n_base * r.SPEED / 40.0)


def same_state(reference, values):
    """Bitwise populations and identical NaN masks; returns a short description or None."""
    for level, (a, b) in enumerate(zip(reference, values)):
        a, b = np.asarray(a), np.asarray(b)
        if a.shape != b.shape:
            return f"level {level}: shape {a.shape} != {b.shape}"
        if not np.array_equal(np.isnan(a), np.isnan(b)):
            return f"level {level}: NaN mask differs"
        finite = np.isfinite(a)
        if a.dtype != b.dtype or not np.array_equal(a[finite].view(np.uint8), b[finite].view(np.uint8)):
            return f"level {level}: finite bits differ"
    return None


def refined_loop(backend, variant, n):
    if variant == "C":
        backend.step(n)
        return
    for _ in range(n):
        backend.step(1)
        if variant == "A":
            backend.read_raw_exchange(0, level=1, window="last_base_step")


def uniform_loop(sim, variant, n):
    if variant == "E'":
        sim.run(n)
        return
    done = 0
    while done < n:
        batch = min(100, n - done)
        sim.run(batch)
        sim.read_force(0)
        done += batch


def time_refined(a, h, sched):
    from contextlib import ExitStack
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper

    warm, n = sched["refined"]
    rows, reference = [], {}
    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, r.BOUNDS, context=context, coarse_boundary="channel", base_block=r.BASE_BLOCK,
            fine_block=r.FINE_BLOCK, fine_objects=(body,), observe_raw_exchange=True))
        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="READY", case="refined")), flush=True)
        for repeat in range(a.repeats):
            for variant in ("A", "B", "C"):
                backend.initialize(*fresh)  # resets step count; same fresh state every time
                refined_loop(backend, variant, warm)
                start = time.perf_counter()
                refined_loop(backend, variant, n)
                seconds = time.perf_counter() - start
                values = [np.array(v, copy=True) for v in backend._download_raw().values]
                mismatch = None
                if repeat == 0 and variant == "A":
                    reference["A"] = values
                else:
                    mismatch = same_state(reference["A"], values)
                rows.append(dict(variant=variant, repeat=repeat, seconds=seconds,
                                 base_steps=n, step_count=backend.step_count,
                                 parity_vs_A=mismatch or "identical", **VARIANTS[variant]))
                if mismatch:
                    raise RuntimeError(f"refined {variant} repeat {repeat}: {mismatch}")
    return rows


def time_uniform(a, h, sched):
    from CelerisLab.simulation import Simulation

    warm, n = sched["uniform"]
    rows, reference = [], {}
    cfgdict, bodydict, _ = r.config(h, 80, "MRT")
    with TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        (tmp / "lbm.json").write_text(json.dumps(cfgdict))
        (tmp / "body.json").write_text(json.dumps(bodydict))
        with Simulation(str(tmp / "lbm.json"), str(tmp / "body.json"), device_id=0) as sim:
            sim.initialize()
            sim.snapshot()  # fresh state; restore keeps clock and streaming parity
            print(json.dumps(dict(status="READY", case="uniform")), flush=True)
            for repeat in range(a.repeats):
                for variant in ("E", "E'"):
                    sim.restore()
                    if sim.stepper.step_count % 2:
                        raise RuntimeError("uniform restore parity changed between variants")
                    uniform_loop(sim, variant, warm)
                    start = time.perf_counter()
                    uniform_loop(sim, variant, n)
                    seconds = time.perf_counter() - start
                    values = [sim.get_ddf()]
                    mismatch = None
                    if repeat == 0 and variant == "E":
                        reference["E"] = values
                    else:
                        mismatch = same_state(reference["E"], values)
                    rows.append(dict(variant=variant, repeat=repeat, seconds=seconds,
                                     uniform_steps=n, parity_vs_E=mismatch or "identical",
                                     **VARIANTS[variant]))
                    if mismatch:
                        raise RuntimeError(f"uniform {variant} repeat {repeat}: {mismatch}")
    return rows


def summarize(rows, tstar_window):
    summary = {}
    for variant in VARIANTS:
        seconds = [row["seconds"] for row in rows if row["variant"] == variant]
        if seconds:
            median = statistics.median(seconds)
            summary[variant] = dict(median_seconds=median, seconds=seconds,
                                    seconds_per_tstar=median / tstar_window,
                                    **VARIANTS[variant])
    if {"A", "B", "C"} <= summary.keys():
        summary["A_minus_B_seconds"] = summary["A"]["median_seconds"] - summary["B"]["median_seconds"]
        summary["B_minus_C_seconds"] = summary["B"]["median_seconds"] - summary["C"]["median_seconds"]
    if {"A", "E"} <= summary.keys():
        summary["A_over_E"] = summary["A"]["median_seconds"] / summary["E"]["median_seconds"]
    if {"C", "E'"} <= summary.keys():
        summary["C_over_Eprime"] = summary["C"]["median_seconds"] / summary["E'"]["median_seconds"]
    summary["interpretation"] = ("Measured comparative differences under GPU variability; "
                                 "not exact causal decompositions or performance claims.")
    return summary


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("--plan", action="store_true")
    a = p.parse_args()
    if a.repeats < 1 or a.base_steps > 20000:
        p.error("repeats must be >= 1 and base-steps <= 20000")
    a.source = a.source.resolve()
    manifest = r.verify_source(a.source, "MRT")
    sched = schedule(a.base_steps, a.warmup_base_steps)
    plan = dict(collision="MRT", manifest_commit=manifest["commit"], schedule=sched,
                repeats=a.repeats, variants=VARIANTS, timer="time.perf_counter wall",
                excluded="construction/JIT, warm-up, state download and parity checks")
    if a.plan:
        print(json.dumps(plan, indent=2))
        return 0
    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)
    h, _ = r.helpers(a.source)
    rows = time_uniform(a, h, sched) + time_refined(a, h, sched)
    result = dict(plan=plan, rows=rows, summary=summarize(rows, sched["tstar_window"]))
    if not all(math.isfinite(row["seconds"]) for row in rows):
        raise RuntimeError("Nonfinite timing")
    (a.output / "timing.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result["summary"], indent=2), flush=True)
    return 0


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