"""Roofline arithmetic and tiled matrix multiplication with traffic counting.

Model used throughout the lesson: two levels of memory. Slow memory is large and has a
limited bandwidth. Fast memory is small, and traffic inside it costs nothing. We count the
elements that cross between the two levels, and the floating-point operations (FLOPs) done.
One multiply-add is 2 FLOPs, the same convention the vendor uses for its peak numbers.

Run `python3 tiled_roofline.py` for the numbers quoted in the lesson, and
`python3 tiled_roofline.py --time` to time numpy and a pure-Python loop on your own machine.
"""
from __future__ import annotations

import sys
import time
from dataclasses import dataclass
from math import ceil, isqrt

import numpy as np

# NVIDIA A100 Tensor Core GPU data sheet (revision "Jun21"), 80GB SXM column.
# BFLOAT16 / FP16 Tensor Core: "312 TFLOPS | 624 TFLOPS*" (* with sparsity). We use the dense value.
# GPU memory bandwidth: 2,039 GB/s (HBM2e). We read GB as 10**9 bytes.
A100_PEAK_FLOPS = 312e12
A100_BANDWIDTH = 2039e9
FP16_BYTES = 2


# --- FLOPs, bytes, intensity, roofline ---------------------------------------------------

def matmul_flops(m: int, n: int, k: int) -> int:
    """C[m, n] = A[m, k] @ B[k, n]: m*n*k multiply-adds, 2 FLOPs each."""
    return 2 * m * n * k


def matmul_min_bytes(m: int, n: int, k: int, elem_bytes: int) -> int:
    """Least traffic possible: read A once, read B once, write C once."""
    return (m * k + k * n + m * n) * elem_bytes


def gemm_intensity(m: int, n: int, k: int, elem_bytes: int) -> float:
    """FLOP per byte if every element moves exactly once (the best case)."""
    return matmul_flops(m, n, k) / matmul_min_bytes(m, n, k, elem_bytes)


def ridge(peak_flops: float, bandwidth: float) -> float:
    """Intensity (FLOP/byte) where the memory slope meets the flat compute roof."""
    return peak_flops / bandwidth


def attainable(intensity: float, peak_flops: float, bandwidth: float) -> float:
    """The roofline: attainable FLOP/s = min(peak, bandwidth * intensity)."""
    return min(peak_flops, bandwidth * intensity)


def classify(intensity: float, peak_flops: float, bandwidth: float) -> str:
    """A tie at the ridge counts as compute-bound."""
    return "memory-bound" if intensity < ridge(peak_flops, bandwidth) else "compute-bound"


def time_floors(flops: float, nbytes: float, peak_flops: float, bandwidth: float) -> tuple[float, float]:
    """The two lower bounds on run time, in seconds: (math alone, traffic alone).

    Neither the math nor the traffic can finish sooner than its own floor, and the two overlap,
    so the run takes at least the larger of the two.
    """
    return flops / peak_flops, nbytes / bandwidth


# --- traffic of a tiled matrix multiply --------------------------------------------------

@dataclass
class Traffic:
    loads: int = 0    # elements read from slow memory
    stores: int = 0   # elements written to slow memory
    flops: int = 0

    def nbytes(self, elem_bytes: int) -> int:
        return (self.loads + self.stores) * elem_bytes

    def intensity(self, elem_bytes: int) -> float:
        return self.flops / self.nbytes(elem_bytes)


def tiled_matmul(A: np.ndarray, B: np.ndarray, b: int) -> tuple[np.ndarray, Traffic]:
    """Multiply with b x b tiles and count what crosses between slow and fast memory.

    For each output tile, the C tile stays in fast memory while we walk along k. Every
    k-step loads one A tile and one B tile. The C tile is written once at the end.
    Edge tiles are smaller when a size is not a multiple of b.
    """
    m, k = A.shape
    k2, n = B.shape
    if k != k2 or b < 1:
        raise ValueError("shapes must agree and the tile size must be at least 1")
    C = np.zeros((m, n), dtype=np.result_type(A, B))
    t = Traffic()
    for i in range(0, m, b):
        for j in range(0, n, b):
            acc = np.zeros((min(b, m - i), min(b, n - j)), dtype=C.dtype)
            for p in range(0, k, b):
                a_tile = A[i:i + b, p:p + b]
                b_tile = B[p:p + b, j:j + b]
                t.loads += a_tile.size + b_tile.size
                t.flops += 2 * a_tile.shape[0] * a_tile.shape[1] * b_tile.shape[1]
                acc += a_tile @ b_tile
            C[i:i + b, j:j + b] = acc
            t.stores += acc.size
    return C, t


def tiled_traffic(m: int, n: int, k: int, b: int) -> Traffic:
    """Closed form of what `tiled_matmul` counts, with no loops over tiles.

    Every A element is loaded once per column of output tiles, ceil(n / b) times.
    Every B element is loaded once per row of output tiles, ceil(m / b) times.
    """
    return Traffic(loads=ceil(n / b) * m * k + ceil(m / b) * k * n,
                   stores=m * n,
                   flops=matmul_flops(m, n, k))


def min_traffic(m: int, n: int, k: int) -> Traffic:
    """Best case: every element of A and B is loaded once, every element of C stored once."""
    return Traffic(loads=m * k + k * n, stores=m * n, flops=matmul_flops(m, n, k))


def naive_traffic(m: int, n: int, k: int) -> Traffic:
    """No reuse: each multiply-add fetches its two operands; C is written once."""
    return Traffic(loads=2 * m * n * k, stores=m * n, flops=matmul_flops(m, n, k))


def tile_steps(m: int, n: int, k: int, b: int) -> list[tuple[int, int, int, int]]:
    """The (output row tile, output column tile, k tile, elements loaded) of each step."""
    steps = []
    for i in range(0, m, b):
        for j in range(0, n, b):
            for p in range(0, k, b):
                bm, bn, bk = min(b, m - i), min(b, n - j), min(b, k - p)
                steps.append((i // b, j // b, p // b, bm * bk + bk * bn))
    return steps


def fast_bytes_needed(b: int, elem_bytes: int) -> int:
    """One A tile, one B tile and one C tile must sit in fast memory together."""
    return 3 * b * b * elem_bytes


def max_tile(fast_bytes: int, elem_bytes: int) -> int:
    """Largest b with 3 * b * b * elem_bytes <= fast_bytes."""
    return isqrt(fast_bytes // (3 * elem_bytes))


def tiled_intensity(n: int, b: int, elem_bytes: int) -> float:
    """FLOP per byte of an n x n x n multiply with b x b tiles: 2n^3 / ((2n^3 / b + n^2) * s)."""
    return tiled_traffic(n, n, n, b).intensity(elem_bytes)


def ridge_tile_if_tile_count_were_fractional(n: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> float:
    """Smallest b with (b / s) / (1 + b / 2n) at or above the ridge, if n / b could be fractional.

    This keeps the term for writing C but drops the whole-tile effect. It is the middle step
    between the rule of thumb s * ridge (306 on the A100) and the exact answer (342 at n = 4096).
    """
    sr = elem_bytes * ridge(peak_flops, bandwidth)
    return sr / (1 - sr / (2 * n))


def smallest_tile_for_ridge(n: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> int:
    """Smallest b whose n x n x n tiled multiply is no longer memory-bound.

    Not the same as the rule of thumb elem_bytes * ridge (306 on the A100), for two causes.
    Writing C adds traffic, which alone moves the answer to 318. The tile count per side,
    ceil(n / b), must also be a whole number: at n = 4096, 13 tiles per side give 151.7 and
    12 give 163.8 (the ridge is 153.0), so b must be at least 4096 / 12 = 341.3, which is 342.
    """
    r = ridge(peak_flops, bandwidth)
    return next(b for b in range(1, n + 1) if tiled_intensity(n, b, elem_bytes) >= r)


# --- decode versus prefill ---------------------------------------------------------------

def batch_to_reach_ridge(n: int, k: int, elem_bytes: int, peak_flops: float, bandwidth: float) -> int:
    """Smallest number of tokens m for which (m x k) @ (k x n) is no longer memory-bound."""
    r = ridge(peak_flops, bandwidth)
    m = 1
    while gemm_intensity(m, n, k, elem_bytes) < r:
        m += 1
    return m


def weight_read_seconds(params: float, elem_bytes: int, bandwidth: float) -> float:
    """Time to read every weight of a model once: the floor for one decoding step."""
    return params * elem_bytes / bandwidth


def decode_tokens_per_s(params: float, elem_bytes: int, bandwidth: float) -> float:
    """Upper bound on one decoding request if each token reads every weight once.

    Ignores the attention cache (stored keys and values) that each token also reads.
    """
    return 1 / weight_read_seconds(params, elem_bytes, bandwidth)


# --- the worked example and a slow reference ---------------------------------------------

EXAMPLE_A = np.array([[1, 2, 0, 1], [3, 1, 2, 0], [0, 1, 1, 2], [2, 0, 1, 1]])
EXAMPLE_B = np.array([[1, 2, 2, 1], [0, 1, 1, 0], [2, 1, 0, 1], [1, 2, 1, 0]])


def naive_matmul(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
    """Three plain Python loops. Slow on purpose: it shows what the library saves you."""
    m, k, n = len(A), len(B), len(B[0])
    C = [[0.0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            s = 0.0
            for p in range(k):
                s += A[i][p] * B[p][j]
            C[i][j] = s
    return C


def time_on_this_machine(n_small: int = 96, n_mm: int = 1024, n_mv: int = 4096) -> dict[str, float]:
    """Time four jobs and return rates (GFLOP/s, plus GB/s for the matrix-vector case). Your numbers depend on your machine."""
    rng = np.random.default_rng(0)

    def best(fn, repeats=3):
        times = []
        for _ in range(repeats):
            t0 = time.perf_counter()
            fn()
            times.append(time.perf_counter() - t0)
        return min(times)

    a_small = rng.standard_normal((n_small, n_small))
    a_list, b_list = a_small.tolist(), a_small.T.tolist()
    A = rng.standard_normal((n_mm, n_mm), dtype=np.float32)
    B = rng.standard_normal((n_mm, n_mm), dtype=np.float32)
    M = rng.standard_normal((n_mv, n_mv), dtype=np.float32)
    x = rng.standard_normal(n_mv, dtype=np.float32)
    t_py = best(lambda: naive_matmul(a_list, b_list), repeats=1)
    t_np = best(lambda: a_small @ a_small.T)
    t_mm = best(lambda: A @ B)
    t_mv = best(lambda: M @ x)
    return {
        f"pure Python loop, {n_small}x{n_small} matmul (GFLOP/s)": matmul_flops(n_small, n_small, n_small) / t_py / 1e9,
        f"numpy, {n_small}x{n_small} matmul (GFLOP/s)": matmul_flops(n_small, n_small, n_small) / t_np / 1e9,
        f"numpy, {n_mm}x{n_mm} matmul, float32 (GFLOP/s)": matmul_flops(n_mm, n_mm, n_mm) / t_mm / 1e9,
        f"numpy, {n_mv}x{n_mv} matrix-vector, float32 (GFLOP/s)": matmul_flops(n_mv, 1, n_mv) / t_mv / 1e9,
        f"numpy, {n_mv}x{n_mv} matrix-vector, matrix bytes read (GB/s)": n_mv * n_mv * 4 / t_mv / 1e9,
    }


# --- report: the numbers quoted in the lesson ---------------------------------------------

def report() -> None:
    s = FP16_BYTES
    r = ridge(A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"A100 SXM 80GB: peak {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, bandwidth {A100_BANDWIDTH / 1e9:.0f} GB/s")
    print(f"ridge point: {r:.1f} FLOP/byte")

    print("\n-- tiny example: 2x2 matmul, float32")
    print("flops", matmul_flops(2, 2, 2), "min bytes", matmul_min_bytes(2, 2, 2, 4),
          "intensity", round(gemm_intensity(2, 2, 2, 4), 4))
    for n in (2, 16, 256, 4096):
        print(f"best-case intensity, n={n}, float32: {gemm_intensity(n, n, n, 4):.3f}")

    n = 4096
    print(f"\n-- naive vs ideal, n={n}, FP16")
    nv = naive_traffic(n, n, n)
    best_t = min_traffic(n, n, n)
    print(f"flops {nv.flops:.4e}")
    print(f"naive bytes {nv.nbytes(s):.4e}, intensity {nv.intensity(s):.3f}")
    print(f"min bytes {best_t.nbytes(s):.4e}, intensity {best_t.intensity(s):.1f}")
    print(f"naive attainable {attainable(nv.intensity(s), A100_PEAK_FLOPS, A100_BANDWIDTH) / 1e12:.3f} TFLOP/s "
          f"= {100 * attainable(nv.intensity(s), A100_PEAK_FLOPS, A100_BANDWIDTH) / A100_PEAK_FLOPS:.2f}% of peak")
    t_math, t_naive = time_floors(nv.flops, nv.nbytes(s), A100_PEAK_FLOPS, A100_BANDWIDTH)
    _, t_min = time_floors(best_t.flops, best_t.nbytes(s), A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"time floor from math {t_math * 1e3:.3f} ms; from naive traffic {t_naive * 1e3:.1f} ms "
          f"({t_naive / t_math:.0f} times the math floor); from minimum traffic {t_min * 1e3:.4f} ms")

    print("\n-- 4x4 with 2x2 tiles")
    C, t = tiled_matmul(EXAMPLE_A, EXAMPLE_B, 2)
    print("loads", t.loads, "stores", t.stores, "flops", t.flops, "naive", naive_traffic(4, 4, 4))
    print("C =", C.tolist())
    print(f"intensity per number moved: naive {naive_traffic(4, 4, 4).intensity(1):.2f}, tiled {t.intensity(1):.2f}")

    print("\n-- simulated loop, n=128, FP16: counted by tiled_matmul vs the formula 2 n^3 / b")
    rng = np.random.default_rng(0)
    A, B = rng.standard_normal((128, 128)), rng.standard_normal((128, 128))
    for b in (4, 8, 16, 32, 64, 128):
        C, t = tiled_matmul(A, B, b)
        assert np.allclose(C, A @ B)
        print(f"b={b}: loads {t.loads}, formula {2 * 128**3 // b}, stores {t.stores}, flops {t.flops}, "
              f"intensity {t.intensity(s):.2f} FLOP/byte")

    print(f"\n-- tile size sweep, n={n}, FP16")
    print("b, loads (elements), traffic/min, intensity, fast bytes needed")
    for b in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 4096):
        t = tiled_traffic(n, n, n, b)
        print(b, t.loads, round(t.nbytes(s) / best_t.nbytes(s), 2), round(t.intensity(s), 2),
              fast_bytes_needed(b, s))

    print("\n-- capacity (example sizes, not a claim about any chip)")
    for kib in (16, 64, 256, 1024):
        b = max_tile(kib * 1024, s)
        print(f"{kib} KiB -> b <= {b}, intensity ~ {tiled_intensity(n, b, s):.1f}")
    rule = s * r
    b_need = smallest_tile_for_ridge(n, s, A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"large-n rule: intensity ~ b / elem_bytes, so the ridge {r:.1f} needs b ~ {rule:.0f}")
    b_frac = ridge_tile_if_tile_count_were_fractional(n, s, A100_PEAK_FLOPS, A100_BANDWIDTH)
    print(f"with the C-store term but a fractional tile count: b >= {b_frac:.1f}, so {ceil(b_frac)}")
    print(f"exact at n={n}: smallest b at or above the ridge is {b_need}; "
          f"intensity there {tiled_intensity(n, b_need, s):.1f}, one smaller {tiled_intensity(n, b_need - 1, s):.1f}; "
          f"fast bytes {fast_bytes_needed(b_need, s)}")

    print("\n-- decode vs prefill: (m x 4096) @ (4096 x 4096), FP16")
    for m in (1, 2, 8, 16, 32, 64, 128, 165, 166, 256, 512, 1024, 4096):
        i = gemm_intensity(m, 4096, 4096, s)
        att = attainable(i, A100_PEAK_FLOPS, A100_BANDWIDTH)
        print(f"m={m}: intensity {i:.2f}, {classify(i, A100_PEAK_FLOPS, A100_BANDWIDTH)}, "
              f"{att / 1e12:.2f} TFLOP/s ({100 * att / A100_PEAK_FLOPS:.1f}% of peak)")
    print("batch to reach ridge:", batch_to_reach_ridge(4096, 4096, s, A100_PEAK_FLOPS, A100_BANDWIDTH))

    params = 7e9
    print(f"\n-- 7B params, FP16: weights {params * s / 1e9:.0f} GB; "
          f"read once = {weight_read_seconds(params, s, A100_BANDWIDTH) * 1e3:.2f} ms; "
          f"at most {decode_tokens_per_s(params, s, A100_BANDWIDTH):.0f} tokens/s at batch 1 (upper bound)")

    print("\n-- check-yourself: 48 KiB of fast memory, 4-byte numbers, n = 4096")
    b = max_tile(48 * 1024, 4)
    print(f"largest tile {b}, intensity {tiled_intensity(n, b, 4):.1f}")


if __name__ == "__main__":
    if "--time" in sys.argv:
        for name, value in time_on_this_machine().items():
            print(f"{name}: {value:.2f}")
    else:
        report()