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  -- runner-like force-read workload
  B  refined: step(1) only, no host reads        -- timing-only, no force history
  C  refined: one step(N) call                   -- timing-only; ledger keeps final base interval only
  E  uniform: sim.run(<=100) + read_force                -- runner-like force-read workload
  E' uniform: one sim.run(2N) call                -- timing-only; at most a whole-window force
                                                    average (different summation order), no history
A/E read forces but discard them: no retained force history, no interval aggregation/validation,
and no macro/probe/CSV cost, so they are not full benchmark throughput.
observe_raw_exchange stays True for A/B/C: B/C omit host reads only.
Status lines: CONSTRUCTED = backend/simulation built (JIT done); WARM_READY = first warm-up and
health check passed for that case.
Parity compares A/B/C with each other and E/E' with each other, never uniform with refined.
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, parity and health checks are outside timed regions. The actual stepping stream is
synchronized explicitly immediately before the timer starts and after the timed loop. Wall time
(perf_counter) is the only timing source; no CUDA-event span is reported as kernel time. An outer
--max-wall-seconds deadline is checked between variants and inside per-step loops.
"""

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=False, force_reads=True,
              force_note="runner-like force-read workload; values discarded, no aggregation"),
    "B": dict(case="refined", force_history_valid=False, timing_only=True,
              force_note="no host reads"),
    "C": dict(case="refined", force_history_valid=False, timing_only=True,
              force_note="raw ledger retains only the final base interval"),
    "E": dict(case="uniform", force_history_valid=False, force_reads=True,
              force_note="runner-like force-read workload; values discarded, no aggregation"),
    "E'": dict(case="uniform", force_history_valid=False, timing_only=True,
               force_note="at most a whole-window accumulated average; different summation order"),
}


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)


# Existing run_refined_dfg safety limits (rho in [.5, 1.5], sampled max Mach <= .2).
RHO_LIMITS, MAX_MACH = (0.5, 1.5), 0.2


def same_state(reference, values):
    """Bitwise populations and identical NaN masks; returns a short description or None.

    Matching NaN is accepted only as a storage mask (inactive/ghost/solid slots); physical
    leaf health is checked separately. Any infinity in either state is a mismatch.
    """
    reference, values = list(reference), list(values)
    if len(reference) != len(values):
        return f"level count {len(reference)} != {len(values)}"
    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 np.isinf(a).any() or np.isinf(b).any():
            return f"level {level}: infinite population"
        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 uniform_populations_finite(ddf, fluid):
    """Simulation.get_ddf() may be flat (9*ny*nx); view it as (Q, cells) before masking."""
    ddf, fluid = np.asarray(ddf), np.asarray(fluid, bool).ravel()
    if fluid.size == 0 or ddf.size % fluid.size:
        raise RuntimeError(f"DDF size {ddf.size} incompatible with {fluid.size} cells")
    return bool(np.isfinite(ddf.reshape(ddf.size // fluid.size, fluid.size)[:, fluid]).all())


def physical_health(fields):
    """fields: (rho, ux, uy) arrays already restricted to physical leaves/fluid."""
    if not fields or any(not np.isfinite(v).all() or v.size == 0 for f in fields for v in f):
        raise RuntimeError("nonfinite or empty physical macroscopic field")
    rho_min = float(min(f[0].min() for f in fields))
    rho_max = float(max(f[0].max() for f in fields))
    max_mach = float(max(np.sqrt(3 * (f[1] ** 2 + f[2] ** 2)).max() for f in fields))
    if not (RHO_LIMITS[0] <= rho_min <= rho_max <= RHO_LIMITS[1] and max_mach <= MAX_MACH):
        raise RuntimeError(f"health limits violated: rho [{rho_min}, {rho_max}], Mach {max_mach}")
    return dict(rho_min=rho_min, rho_max=rho_max, max_mach=max_mach)


def refined_health(backend, masks):
    arrays = backend.download()  # physical FP64 populations; non-leaves NaN by contract
    if any(not np.isfinite(g[:, m]).all() for g, m in zip(arrays, masks)):
        raise RuntimeError("nonfinite physical leaf population")
    return physical_health([tuple(v[m] for v in r.macros(g)) for g, m in zip(arrays, masks)])


def uniform_health(sim, fluid):
    m = sim.get_macroscopic()
    return physical_health([tuple(np.asarray(m[k], np.float64).reshape(fluid.shape)[fluid]
                                  for k in ("rho", "ux", "uy"))])


def check_deadline(deadline):
    if time.monotonic() >= deadline:
        raise TimeoutError("timing harness exceeded --max-wall-seconds")


def sync_refined(backend):
    with backend.ctx.activate():
        backend._stream.synchronize()


def sync_uniform(sim):
    with sim.ctx.activate():
        sim.stream.synchronize()


def refined_loop(backend, variant, n, deadline):
    if variant == "C":
        backend.step(n)
        return
    for i in range(n):
        backend.step(1)
        if variant == "A":
            backend.read_raw_exchange(0, level=1, window="last_base_step")
        if i % 100 == 99:
            check_deadline(deadline)


def uniform_loop(sim, variant, n, deadline):
    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
        check_deadline(deadline)


def time_refined(a, h, sched):
    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

    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))
        masks = [channel_masks(pair.coarse_shape)[0] & ~pair.covered,  # as run_refined_dfg
                 (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 ("A", "B", "C"):
                check_deadline(a.deadline)
                backend.initialize(*fresh)  # resets step count; same fresh state every time
                refined_loop(backend, variant, warm, a.deadline)
                warm_health = refined_health(backend, masks)  # untimed
                if repeat == 0 and variant == "A":
                    print(json.dumps(dict(status="WARM_READY", case="refined")), flush=True)
                sync_refined(backend)
                start = time.perf_counter()
                refined_loop(backend, variant, n, a.deadline)
                sync_refined(backend)
                seconds = time.perf_counter() - start
                end_health = refined_health(backend, masks)  # untimed
                values = [np.array(v, copy=True) for v in backend._download_raw().values]
                if not all(np.isfinite(v[:, m]).all() for v, m in zip(values, backend._state)):
                    raise RuntimeError(f"refined {variant} repeat {repeat}: nonfinite leaf populations")
                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,
                                 health_after_warmup=warm_health, health_after_window=end_health,
                                 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
    from CelerisLab.lbm.descriptors import FLUID

    warm, n = sched["uniform"]
    rows, reference = [], {}
    cfgdict, bodydict, params = 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
            fluid = (sim.get_flags() & FLUID) != 0
            ny, nx = params["ny"], params["nx"]
            fluid = fluid.reshape(ny, nx)
            print(json.dumps(dict(status="CONSTRUCTED", case="uniform")), flush=True)
            for repeat in range(a.repeats):
                for variant in ("E", "E'"):
                    check_deadline(a.deadline)
                    sim.restore()
                    if sim.stepper.step_count % 2:
                        raise RuntimeError("uniform restore parity changed between variants")
                    uniform_loop(sim, variant, warm, a.deadline)
                    warm_health = uniform_health(sim, fluid)  # untimed
                    if repeat == 0 and variant == "E":
                        print(json.dumps(dict(status="WARM_READY", case="uniform")), flush=True)
                    sync_uniform(sim)
                    start = time.perf_counter()
                    uniform_loop(sim, variant, n, a.deadline)
                    sync_uniform(sim)
                    seconds = time.perf_counter() - start
                    end_health = uniform_health(sim, fluid)  # untimed
                    values = [sim.get_ddf()]
                    if not uniform_populations_finite(values[0], fluid):
                        raise RuntimeError(f"uniform {variant} repeat {repeat}: nonfinite fluid populations")
                    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",
                                     health_after_warmup=warm_health, health_after_window=end_health,
                                     **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("--max-wall-seconds", type=float, help="Required outer timeout for GPU runs")
    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 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")
    a.deadline = time.monotonic() + a.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)
    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())