code/experiments/exp_fig3_reconstruction.py

"""Claim 6(a) — Figure 3 reproduction: FP32 reconstruction error by exponent.

Protocol matches paper §4.4: evaluate weight-splitting schemes EXHAUSTIVELY over
all finite FP32 bitstrings (4,278,190,082 finite values incl. ±0; zeros excluded
from relative-error means), grouped by the value's biased exponent field.

Methods (target dtype bf16 and fp16):
  none      : plain downcast, no correction            (Fig. 3 "no error correction")
  float     : error stored in same low-precision float (Zamirai-style, "BF16+BF16")
  ulp_int8  : paper's ULP-normalized INT8 correction   ("ours, 24-bit")
  ulp_int16 : paper's ULP-normalized INT16 correction  ("ours, 32-bit")

Outputs per (target, method, exponent): mean/max relative error, bitwise-exact
fraction, nonfinite-reconstruction count. Paper stats checked:
  - BF16+INT16 mean rel err < 1e-9 (normal range) and 99.92% bitwise-exact overall
  - BF16+BF16 error > 1e-6, comparable to 24-bit ULP format
  - FP16+INT16 exact in fp16 normal range; 24-bit worst-case < 1e-6 vs 1e-4
"""

import json
import os
import sys
import time

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from flashsim import splitting as S  # noqa: E402

CHUNK = 2**24
N_EXP = 255  # biased exponent field 0..254 (255 = inf/nan, excluded)

METHODS = {
    "none": lambda x, lp: S.roundtrip_no_correction(x, lp),
    "float": lambda x, lp: S.roundtrip_float_error(x, lp),
    "ulp_int8": lambda x, lp: S.roundtrip(x, lp, torch.int8),
    "ulp_int16": lambda x, lp: S.roundtrip(x, lp, torch.int16),
}
TARGETS = {"bf16": torch.bfloat16, "fp16": torch.float16}


def main(fast: bool = False):
    acc = {t: {m: {"err_sum": torch.zeros(N_EXP, dtype=torch.float64),
                   "err_max": torch.zeros(N_EXP, dtype=torch.float64),
                   "exact": torch.zeros(N_EXP, dtype=torch.float64),
                   "nonfinite": torch.zeros(N_EXP, dtype=torch.float64)}
               for m in METHODS} for t in TARGETS}
    counts = torch.zeros(N_EXP, dtype=torch.float64)
    n_total = n_zero = 0
    step = 64 if fast else 1  # fast mode: every 64th chunk (stratified smoke)
    t0 = time.time()
    chunk_ids = range(0, 2**32 // CHUNK, step)
    for ci, c in enumerate(chunk_ids):
        bits = (torch.arange(CHUNK, dtype=torch.int64) + c * CHUNK).to(torch.uint32)
        x = bits.view(torch.int32).view(torch.float32)
        finite = torch.isfinite(x)
        x = x[finite]
        if x.numel() == 0:
            continue
        n_total += x.numel()
        absx = x.abs()
        nz = absx > 0
        n_zero += int((~nz).sum())
        exp_field = ((x.view(torch.int32).to(torch.int64) >> 23) & 0xFF)
        counts += torch.bincount(exp_field[nz], minlength=N_EXP).to(torch.float64)
        for tname, lp in TARGETS.items():
            for mname, fn in METHODS.items():
                xh = fn(x, lp)
                a = acc[tname][mname]
                # bitwise-exact fraction over all finite values (incl. zeros)
                exact = (xh.view(torch.int32) == x.view(torch.int32))
                a["exact"] += torch.bincount(exp_field[exact], minlength=N_EXP).to(torch.float64)
                fin = torch.isfinite(xh)
                a["nonfinite"] += torch.bincount(exp_field[~fin], minlength=N_EXP).to(torch.float64)
                sel = nz & fin
                rel = ((xh[sel] - x[sel]).abs().double() / absx[sel].double())
                e = exp_field[sel]
                a["err_sum"] += torch.bincount(e, weights=rel, minlength=N_EXP)
                m = torch.zeros(N_EXP, dtype=torch.float64)
                m.scatter_reduce_(0, e, rel, reduce="amax")
                a["err_max"] = torch.maximum(a["err_max"], m)
        if ci % 16 == 0:
            done = (ci + 1) / len(list(chunk_ids))
            print(f"[{time.time()-t0:7.1f}s] chunk {ci+1}/{len(list(chunk_ids))} ({100*done:.1f}%)", flush=True)

    # exact-count denominator per exponent (finite values incl zeros): counts
    # holds nonzero-only; zeros live in exponent 0. Track separately:
    zeros_per_exp = torch.zeros(N_EXP, dtype=torch.float64)
    zeros_per_exp[0] = n_zero
    denom_all = counts + zeros_per_exp

    rows = []
    for tname in TARGETS:
        for mname in METHODS:
            a = acc[tname][mname]
            mean = torch.where(counts > 0, a["err_sum"] / counts.clamp_min(1), torch.zeros(N_EXP, dtype=torch.float64))
            for e in range(N_EXP):
                if denom_all[e] == 0:
                    continue
                rows.append(dict(target=tname, method=mname, exp_biased=e,
                                 exp_unbiased=e - 127 if e > 0 else -126,
                                 n=int(counts[e]),
                                 mean_rel_err=float(mean[e]),
                                 max_rel_err=float(a["err_max"][e]),
                                 exact_frac=float(a["exact"][e] / denom_all[e]),
                                 nonfinite=int(a["nonfinite"][e])))
    os.makedirs("repro_flashoptim/outputs", exist_ok=True)
    import csv
    with open("repro_flashoptim/outputs/fig3_reconstruction.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader(); w.writerows(rows)

    # headline stats vs paper
    def overall(t, m, field="exact"):
        a = acc[t][m]
        return float(a[field].sum() / denom_all.sum())
    # mean rel err over NORMAL range only (biased exp 1..254 for fp32... but
    # restrict to bf16-normal magnitudes for the bf16 headline: exp >= 1)
    def mean_err_normal(t, m, lo=1, hi=254):
        a = acc[t][m]
        c = counts[lo:hi + 1].sum()
        return float(a["err_sum"][lo:hi + 1].sum() / c)

    def exact_frac_range(t, m, lo, hi):
        a = acc[t][m]
        return float(a["exact"][lo:hi + 1].sum() / denom_all[lo:hi + 1].sum())
    FP16_LO, FP16_HI = 113, 142  # fp32 biased exponents inside fp16 normal range
    summary = {
        "n_finite_values_swept": n_total,
        "exhaustive": not fast,
        "bf16_ulp_int16_exact_frac_overall": overall("bf16", "ulp_int16"),
        "paper_bf16_int16_exact": 0.9992,
        "bf16_ulp_int16_mean_rel_err_normal": mean_err_normal("bf16", "ulp_int16"),
        "paper_bf16_int16_err_bound": 1e-9,
        "bf16_float_mean_rel_err_normal": mean_err_normal("bf16", "float"),
        "paper_bf16_bf16_err_bound_exceeds": 1e-6,
        "bf16_ulp_int8_mean_rel_err_normal": mean_err_normal("bf16", "ulp_int8"),
        "bf16_none_mean_rel_err_normal": mean_err_normal("bf16", "none"),
        "fp16_ulp_int16_exact_frac_fp16normal": exact_frac_range("fp16", "ulp_int16", FP16_LO, FP16_HI),
        "fp16_ulp_int16_mean_rel_err_fp16normal": mean_err_normal("fp16", "ulp_int16", FP16_LO, FP16_HI),
        "fp16_float_mean_rel_err_fp16normal": mean_err_normal("fp16", "float", FP16_LO, FP16_HI),
        "fp16_ulp_int8_mean_rel_err_fp16normal": mean_err_normal("fp16", "ulp_int8", FP16_LO, FP16_HI),
        "fp16_ulp_int8_max_rel_err_fp16normal": float(acc["fp16"]["ulp_int8"]["err_max"][FP16_LO:FP16_HI + 1].max()),
        "paper_fp16_24bit_worstcase_under": 1e-6,
        "elapsed_s": round(time.time() - t0, 1),
    }
    with open("repro_flashoptim/outputs/fig3_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main(fast="--fast" in sys.argv)