#!/usr/bin/env python3
"""Hands-on P14 — hardware-aware ML, assembled from seven lego blocks."""
import time
import numpy as np
from _harness import block, run_all

def best(fn, reps):
    out = []
    for _ in range(reps):
        t0 = time.perf_counter(); fn(); out.append(time.perf_counter() - t0)
    return min(out)

@block(1, "Measure the machine, not the spec sheet", "the roofline is per-dtype, and that bites")
def b1(s, show):
    n = 1024
    peaks = {}
    for dt in (np.float64, np.float32):
        A = np.random.rand(n, n).astype(dt); B = np.random.rand(n, n).astype(dt)
        peaks[np.dtype(dt).name] = 2 * n**3 / best(lambda: A @ B, 5)
    # peak bandwidth: a streaming triad, far larger than any cache
    N = 24_000_000
    x = np.random.rand(N); y = np.random.rand(N); z = np.empty(N)
    def triad(): np.multiply(x, 2.0, out=z); np.add(z, y, out=z)
    tb = best(triad, 5)
    bw = 3 * N * 8 / tb
    if show:
        print(f"  peak bandwidth : {bw/1e9:>8.1f} GB/s   "
              f"(triad over {N*8/1e6:.0f} MB in {tb*1000:.1f} ms)")
        print(f"  {'dtype':<10}{'peak GFLOP/s':>15}{'ridge point':>15}")
        for k, v in peaks.items():
            print(f"  {k:<10}{v/1e9:>15.1f}{v/bw:>13.2f}")
        print(f"  fp32 is {peaks['float32']/peaks['float64']:.2f}x fp64 -- two SIMD lanes")
        print("  per register instead of one, and on this machine a little more than")
        print("  the theoretical 2x because the fp32 kernel also gets better cache")
        print("  reuse per byte.\n")
        print("  THERE IS NO SUCH THING AS 'the' ROOFLINE. There is one per dtype, and")
        print(f"  the ridge moves with it: {peaks['float64']/bw:.1f} FLOP/byte in fp64, "
              f"{peaks['float32']/bw:.1f} in fp32. Mixing")
        print("  them is not a rounding error -- the first version of this file")
        print("  measured peak in fp64 and predicted an fp32 workload, and the")
        print("  assembly reported kernels running at 2x the speed of light. A ratio")
        print("  below 1.0 against a roofline is never a fast kernel; it is always a")
        print("  broken model, and it is the most useful bug the roofline can produce")
        print("  because it is impossible to rationalise away. There are exactly two")
        print("  ways to get one: the wrong ceiling (this bug) or the wrong byte count")
        print("  (a working set that never left cache -- the assembly measures that")
        print("  one too).")
        print("  Both numbers are MEASURED. The vendor's peak assumes an FMA on every")
        print("  port every cycle, which no real kernel reaches; numbers.md records")
        print("  this machine's bandwidth at 57.5-99.1 GB/s depending on working set,")
        print("  and this triad sits at the streaming end of that range.")
    return {"flops": peaks["float64"], "flops32": peaks["float32"], "bw": bw,
            "ridge": peaks["float64"] / bw, "ridge32": peaks["float32"] / bw}

@block(2, "Arithmetic intensity", "the property that decides which wall you hit")
def b2(s, show):
    if show:
        print(f"  fp64 ridge = {s['ridge']:.2f} FLOP/byte on this machine "
              f"(fp32 ridge = {s['ridge32']:.2f})\n")
        print(f"  {'kernel':<28}{'FLOPs':>12}{'bytes':>12}{'intensity':>12}{'bound by':>11}")
        N = 4_000_000
        cases = [
            ("vector add (a+b)",        N,        3*N*8),
            ("scale (a*2)",             N,        2*N*8),
            ("dot product",            2*N,       2*N*8),
            ("matmul 64x64",       2*64**3,    3*64*64*8),
            ("matmul 1024x1024", 2*1024**3, 3*1024*1024*8),
            ("attention, seq 1024",  4*1024**2*64, 3*1024*64*8),
        ]
        for lbl, f, b in cases:
            ai = f / b
            print(f"  {lbl:<28}{f/1e6:>10.1f}M{b/1e6:>10.1f}M{ai:>12.2f}"
                  f"{('COMPUTE' if ai > s['ridge'] else 'MEMORY'):>11}")
        print("  Elementwise operations have intensity below 1 and can never be")
        print("  compute-bound on any machine built since about 1990 -- the ridge has")
        print("  been climbing for thirty years while DRAM latency barely moved.")
        print("  Matmul's intensity grows as O(n), which is the entire reason deep")
        print("  learning runs on hardware designed for it.")
    return {}

@block(3, "Predict, then measure", "a roofline is a falsifiable claim")
def b3(s, show):
    if show:
        print(f"  {'kernel':<26}{'intensity':>11}{'predicted':>13}{'measured':>12}"
              f"{'ratio':>8}")
        N = 8_000_000
        a = np.random.rand(N); b = np.random.rand(N); c = np.empty(N)
        tests = []
        t = best(lambda: np.add(a, b, out=c), 5)
        tests.append(("vector add", N, 3*N*8, t))
        t = best(lambda: np.multiply(a, 2.0, out=c), 5)
        tests.append(("scale by constant", N, 2*N*8, t))
        t = best(lambda: float(a @ b), 5)
        tests.append(("dot product", 2*N, 2*N*8, t))
        for n in (256, 1024):
            A = np.random.rand(n, n); B = np.random.rand(n, n)
            t = best(lambda: A @ B, 5)
            tests.append((f"matmul {n}x{n}", 2*n**3, 3*n*n*8, t))
        for lbl, f, byt, t in tests:
            ai = f / byt
            pred = min(s["flops"], ai * s["bw"])          # the roofline itself
            meas = f / t
            print(f"  {lbl:<26}{ai:>11.2f}{pred/1e9:>11.1f}G{meas/1e9:>10.1f}G"
                  f"{meas/pred:>8.2f}")
        print("  A ratio near 1.0 means the roofline explained the kernel. Below 1.0")
        print("  means something else is the limit -- latency, a missing")
        print("  vectorisation, an unaligned access. Above 1.0 means a modelling")
        print("  error: usually the kernel hit cache and never touched DRAM, so the")
        print("  'bytes' figure is fiction. The model earns trust by being wrong in")
        print("  ways you can explain.")
    return {}

@block(4, "Tiling", "the same FLOPs, a different traffic pattern")
def b4(s, show):
    def naive(A, B):
        n = A.shape[0]; C = np.zeros((n, n))
        for i in range(n):
            for k in range(n):
                C[i] += A[i, k] * B[k]
        return C
    def tiled(A, B, T=64):
        n = A.shape[0]; C = np.zeros((n, n))
        for i0 in range(0, n, T):
            for k0 in range(0, n, T):
                for j0 in range(0, n, T):
                    C[i0:i0+T, j0:j0+T] += (A[i0:i0+T, k0:k0+T]
                                            @ B[k0:k0+T, j0:j0+T])
        return C
    if show:
        n = 512
        A = np.random.rand(n, n); B = np.random.rand(n, n)
        ref = A @ B
        t_bl = best(lambda: A @ B, 3)
        t_na = best(lambda: naive(A, B), 1)
        rows = [("row-at-a-time (python loop)", t_na, naive(A, B))]
        for T in (32, 64, 128, 256):
            tt = best(lambda: tiled(A, B, T), 3)
            rows.append((f"tiled, T={T}", tt, tiled(A, B, T)))
        rows.append(("numpy (BLAS)", t_bl, ref))
        print(f"  C = A @ B, {n}x{n}, {2*n**3/1e9:.2f} GFLOP")
        print(f"  {'implementation':<30}{'time':>10}{'GFLOP/s':>10}"
              f"{'vs BLAS':>10}{'correct':>9}")
        for lbl, tt, out in rows:
            print(f"  {lbl:<30}{tt*1000:>8.1f}ms{2*n**3/tt/1e9:>10.1f}"
                  f"{t_bl/tt:>9.2f}x"
                  f"{str(bool(np.allclose(out, ref))):>9}")
        wss = lambda T: 3 * T * T * 8 / 1024
        print(f"  working set per tile: T=32 -> {wss(32):.0f} KB, "
              f"T=64 -> {wss(64):.0f} KB, T=128 -> {wss(128):.0f} KB, "
              f"T=256 -> {wss(256):.0f} KB")
        print("  The tiles here still call BLAS, so this measures BLOCKING, not")
        print("  hand-written inner loops: how much you lose by cutting a big matmul")
        print("  into small ones. The loss is real and it comes from per-call")
        print("  overhead plus reduced reuse -- which is the same trade the tile-size")
        print("  choice makes inside a real GEMM, one level down.")
    return {"tiled": tiled}

@block(5, "Quantisation", "fewer bytes per weight IS higher arithmetic intensity")
def b5(s, show):
    if show:
        rng = np.random.default_rng(5)
        n = 2048
        W = rng.normal(0, 0.5, (n, n)).astype(np.float32)
        x = rng.normal(0, 1, n).astype(np.float32)
        scale = np.abs(W).max() / 127.0
        Wq = np.clip(np.round(W / scale), -127, 127).astype(np.int8)
        ref = W @ x
        deq = (Wq.astype(np.float32) * scale) @ x
        t32 = best(lambda: W @ x, 20)
        t8 = best(lambda: (Wq.astype(np.float32) * scale) @ x, 20)
        t8b = best(lambda: Wq.astype(np.float32) @ x, 20)
        err = np.abs(deq - ref).max() / np.abs(ref).max()
        cos = float(deq @ ref / (np.linalg.norm(deq) * np.linalg.norm(ref)))
        print(f"  {n}x{n} weight matrix, matrix-vector product (the decode shape)")
        print(f"  {'precision':<20}{'weight bytes':>14}{'time':>10}{'GB/s':>9}"
              f"{'rel error':>12}")
        for lbl, byts, tt, e in (("float32", W.nbytes, t32, 0.0),
                                 ("int8 + dequant", Wq.nbytes, t8, err)):
            print(f"  {lbl:<20}{byts/1e6:>12.1f}M{tt*1e6:>9.0f}us"
                  f"{byts/tt/1e9:>9.1f}{e:>12.2e}")
        print(f"  cosine similarity of the two outputs: {cos:.6f}")
        print(f"  4x fewer weight bytes; measured speedup {t32/t8:.2f}x, and the")
        print("  dequantisation itself costs most of what the smaller load saved.")
        print("  A real int8 kernel keeps the arithmetic in int8 and dequantises the")
        print("  ACCUMULATOR once, which is why production quantisation needs kernel")
        print("  support and not just a smaller dtype in memory. Numpy has no int8")
        print("  GEMM, so what this block honestly measures is the memory saving and")
        print("  the accuracy cost -- both real, and the speedup is the part you")
        print("  cannot get without writing the kernel.")
    return {}

@block(6, "Batching and Little's Law", "the only free speedup in a memory-bound regime")
def b6(s, show):
    if show:
        rng = np.random.default_rng(6)
        n = 2048
        W = rng.normal(0, .5, (n, n)).astype(np.float32)
        print(f"  one {n}x{n} weight matrix, batch of B vectors:")
        print(f"  {'batch':>7}{'time':>10}{'per-item':>11}{'GFLOP/s':>10}"
              f"{'intensity':>11}{'weight reads':>14}")
        t1 = None
        for B in (1, 2, 8, 32, 128):
            X = rng.normal(0, 1, (n, B)).astype(np.float32)
            t = best(lambda: W @ X, 10)
            t1 = t1 or t
            f = 2 * n * n * B
            byt = W.nbytes + X.nbytes + 4 * n * B
            print(f"  {B:>7}{t*1e6:>8.0f}us{t*1e6/B:>10.1f}us"
                  f"{f/t/1e9:>10.1f}{f/byt:>11.2f}{'1':>14}")
        print("  The weight matrix is read ONCE regardless of batch size, so every")
        print("  extra request in the batch is nearly free until the kernel becomes")
        print("  compute-bound. That is why LLM serving batches aggressively and why")
        print("  batch-1 latency is the worst possible operating point: you pay the")
        print("  full 16 MB weight read to produce a single token.")
        print("  Little's Law gives the other half: L = lambda x W. To keep a batch of")
        print("  32 in flight at 20 ms per batch you need 1600 requests/second of")
        print("  arrival. Below that, the batch never fills and you are choosing")
        print("  between latency and utilisation -- which is what a scheduler's")
        print("  max-wait parameter actually configures.")
    return {}

@block(7, "The decode wall", "why generation is memory-bound and prefill is not")
def b7(s, show):
    if show:
        print("  A transformer layer, hidden d, batch B, sequence S. Weights are")
        print("  ~12d^2 bytes in fp16; the matmuls are ~24 B S d^2 FLOPs.")
        print(f"  {'phase':<22}{'B':>4}{'S':>7}{'intensity':>12}{'bound by':>11}"
              f"{'note':>22}")
        d = 4096
        for lbl, B, S in (("prefill, 2k prompt", 1, 2048), ("decode, 1 token", 1, 1),
                          ("decode, batch 32", 32, 1), ("decode, batch 256", 256, 1)):
            flops = 24 * B * S * d * d
            byts = 12 * d * d + 4 * B * S * d
            ai = flops / byts
            note = "reads 200MB per token" if B == 1 and S == 1 else ""
            print(f"  {lbl:<22}{B:>4}{S:>7}{ai:>12.1f}"
                  f"{('COMPUTE' if ai > s['ridge32'] else 'MEMORY'):>11}{note:>22}")
        print(f"  (fp32 ridge on this machine = {s['ridge32']:.1f} FLOP/byte; on an H100")
        print("  with ~990 TFLOP/s and ~3.35 TB/s it is about 295, so the same table")
        print("  on a GPU pushes even batch-256 decode into the memory-bound column.)")
        print("  Prefill has S=2048 tokens sharing one weight read, so it is compute-")
        print("  bound and scales with FLOPs. Decode has S=1: the SAME weights are")
        print("  read to produce a single token. Batching is the only lever that")
        print("  raises decode intensity, which is the whole reason continuous")
        print("  batching, paged attention and speculative decoding exist -- all")
        print("  three are attempts to get more work per weight read. See proofs.md")
        print("  P6 for the derivation.")
    return {}

def assembly(s):
    print("\nSeven blocks = a performance model. Predict a workload before running it.\n")
    rng = np.random.default_rng(14)
    d, L = 1024, 6
    Ws = [rng.normal(0, .02, (d, d)).astype(np.float32) for _ in range(L)]
    def layer(x, W): return np.maximum(x @ W, 0)
    print(f"  Workload: a {L}-layer MLP, width {d}, fp32, batch B.")
    print(f"  Per batch: {L} matmuls = {2*L*d*d/1e6:.1f} MFLOP per item,")
    print(f"  weights = {L*d*d*4/1e6:.1f} MB read once per batch.\n")
    print(f"  {'batch':>6}{'intensity':>11}{'roofline says':>15}{'predicted':>12}"
          f"{'measured':>11}{'ratio':>8}{'bound':>9}")
    ratio_b1 = None
    for B in (1, 4, 16, 64, 256):
        X = rng.normal(0, 1, (B, d)).astype(np.float32)
        flops = 2 * L * B * d * d
        byts = L * d * d * 4 + 2 * B * d * 4 * L
        ai = flops / byts
        pred_rate = min(s["flops32"], ai * s["bw"])      # fp32 workload -> fp32 peak
        pred_t = flops / pred_rate
        def run():
            x = X
            for W in Ws: x = layer(x, W)
            return x
        t = best(run, 5)
        bound = "COMPUTE" if ai > s["ridge32"] else "MEMORY"
        if B == 1: ratio_b1 = t / pred_t
        print(f"  {B:>6}{ai:>11.2f}{bound:>15}{pred_t*1e6:>10.0f}us"
              f"{t*1e6:>9.0f}us{t/pred_t:>8.2f}{bound:>9}")
    print("\n  Now hold the batch at 1 and grow the weights instead, so the same")
    print("  model is priced against working sets that do and do not fit in cache:\n")
    print(f"  {'width':>7}{'layers':>8}{'weight MB':>12}{'predicted':>12}"
          f"{'measured':>11}{'ratio':>8}")
    fp_ratios = []
    for dd, LL in ((1024, 6), (2048, 6), (2048, 12), (4096, 8)):
        Ws2 = [rng.normal(0, .02, (dd, dd)).astype(np.float32) for _ in range(LL)]
        X = rng.normal(0, 1, (1, dd)).astype(np.float32)
        byts = LL * dd * dd * 4 + 2 * dd * 4 * LL
        def run2():
            v = X
            for W in Ws2: v = np.maximum(v @ W, 0)
            return v
        tt = best(run2, 5); pred = byts / s["bw"]
        fp_ratios.append(tt / pred)
        print(f"  {dd:>7}{LL:>8}{LL*dd*dd*4/1e6:>12.1f}{pred*1e6:>10.0f}us"
              f"{tt*1e6:>9.0f}us{tt/pred:>8.2f}")
    print(f"\n  Look at the 25 MB row twice. The batch table measured it at "
          f"{ratio_b1:.2f} and")
    print(f"  the footprint table at {fp_ratios[0]:.2f} -- the same weights, the same")
    print("  arithmetic, both sitting AT the bound rather than comfortably above it,")
    print(f"  and wandering by {abs(ratio_b1-fp_ratios[0]):.2f} between two runs in the same process.")
    print("  A ratio that hovers at or below 1.0 is the signature of a PARTLY cache-")
    print("  resident working set: the weights are re-read every iteration, some")
    print("  fraction survives in this machine's last-level cache, and how large that")
    print("  fraction is depends on what else touched memory first. The model charged")
    print("  the kernel for 25 MB of DRAM traffic that it only partly paid.")
    print(f"  Past 100 MB the ambiguity disappears: {fp_ratios[1]:.2f}, {fp_ratios[2]:.2f}, "
          f"{fp_ratios[3]:.2f}. The working")
    print("  set no longer fits, every byte really does come from DRAM, and the")
    print("  roofline becomes a bound the kernel reaches about half of. The cliff")
    print("  between those two regimes is the cache, measured without ever naming")
    print("  its size -- and it is the same cliff as P12's page-fault curve and")
    print("  P04's Bloom filter, one level up the hierarchy.")
    print("\n  Two different bugs produced a sub-1.0 ratio in this file: an fp64")
    print("  ceiling on an fp32 workload (block 1, a factor of 4) and a byte count")
    print("  that assumed DRAM for data sitting in cache (above). Both were invisible in the")
    print("  absolute timings and both were obvious the moment the number was divided")
    print("  by a bound it could not legally cross. That is what the model is FOR --")
    print("  not predicting runtime, but making a specific class of mistake loud.")
    print("\n  The prediction uses two numbers measured in block 1 and a FLOP count")
    print("  done on paper. No profiler, no counters. Read the batch table's gap")
    print("  where it is largest: at batch 256 we reach roughly half of peak, because")
    print("  a 1024x256 matmul is skinnier than the square one that set the ceiling")
    print("  and there is a full-size ReLU pass between every layer.")
    print("\n  This is the deliverable of hardware-aware ML: not a faster kernel, but")
    print("  the ability to say IN ADVANCE which optimisations can possibly help.")
    print("  If a kernel sits at 0.9 of its roofline, rewriting the inner loop is")
    print("  wasted work and the only remaining moves are algorithmic -- fewer bytes")
    print("  (quantisation, block 5) or more work per byte (batching, block 6).")
    print("\n  Built: measured roofline -> arithmetic intensity -> prediction vs")
    print("  measurement -> tiling -> quantisation -> batching -> the decode wall.")
    print("  Missing, on the project page: hardware counters via perf (m3), a real")
    print("  hand-written GEMM microkernel with register blocking (m5), operator")
    print("  fusion measured end to end (m7), GPU occupancy and warp scheduling")
    print("  (m9-m10), and E2 -- the experiment that finds this machine's cache")
    print("  hierarchy from a latency curve rather than from a spec sheet.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P14 — Hardware-aware ML, block by block")