AI Compiler Drills：Triton + Luminal 双项目 contributor 训练

"""K1 · GEMM tile & resource model — your implementation (run ./check.sh k1).

Pure Python. Given a GEMM shape and a tile configuration (what Triton calls
BLOCK_M / BLOCK_N / BLOCK_K / num_warps / num_stages), predict *before running anything*:

  * how many FLOPs and how many bytes the problem needs at minimum,
  * how many bytes the tiled kernel really moves (A and B are re-read once per tile column/row),
  * how much shared memory and roughly how many registers one CTA needs,
  * how many CTAs fit on one SM (occupancy) and which resource limits it,
  * how many waves of CTAs the grid needs and how full the last wave is,
  * a roofline estimate: is this config memory-, compute- or launch-bound?

These are exactly the numbers you compare against Nsight Compute in T3, and the numbers a
compiler's cost model (Luminal's search, Triton's autotuner pruning) reasons about.
"""
from dataclasses import dataclass
from math import ceil


@dataclass(frozen=True)
class GPU:
    name: str
    sm_count: int
    smem_per_sm: int            # bytes of shared memory per SM
    smem_per_block: int         # max bytes one CTA may use
    regs_per_sm: int            # 32-bit registers per SM
    max_threads_per_sm: int
    max_blocks_per_sm: int
    max_regs_per_thread: int
    peak_flops: float           # FLOP/s for the dtype being modelled (e.g. FP16 Tensor Core)
    peak_bw: float              # bytes/s of DRAM bandwidth
    launch_overhead: float = 5e-6   # seconds; below ~this a kernel is launch-bound


# Approximate public specs. RTX 3070 = GA104, SM86 (the roadmap's starting GPU).
RTX3070 = GPU("RTX 3070 (SM86)", sm_count=46, smem_per_sm=100 * 1024, smem_per_block=99 * 1024,
              regs_per_sm=65536, max_threads_per_sm=1536, max_blocks_per_sm=16, max_regs_per_thread=255,
              peak_flops=40e12, peak_bw=448e9)


@dataclass(frozen=True)
class TileConfig:
    BM: int
    BN: int
    BK: int
    num_warps: int = 4
    num_stages: int = 3


def gemm_flops(M, N, K):
    """One multiply + one add per (m, n, k)."""
    raise NotImplementedError("gemm_flops")


def gemm_min_bytes(M, N, K, in_bytes=2, out_bytes=2):
    """Read A and B once, write C once: the lower bound on DRAM traffic."""
    raise NotImplementedError("gemm_min_bytes")


def arithmetic_intensity(flops, nbytes):
    raise NotImplementedError("arithmetic_intensity")


def tiled_bytes(M, N, K, cfg: TileConfig, in_bytes=2, out_bytes=2):
    """Bytes a tiled kernel loads from DRAM (ignoring L2 hits).

    Each CTA owns a BM x BN tile of C and walks all of K, so the whole of A is read once per
    column of tiles (ceil(N / BN) times) and the whole of B once per row of tiles."""
    raise NotImplementedError("tiled_bytes")


def smem_bytes(cfg: TileConfig, in_bytes=2):
    """num_stages copies of one A tile (BM x BK) and one B tile (BK x BN)."""
    raise NotImplementedError("smem_bytes")


def threads(cfg: TileConfig):
    raise NotImplementedError("threads")


def regs_per_thread(cfg: TileConfig, overhead=40):
    """Rough estimate: the fp32 accumulator tile spread over the CTA's threads, plus a fixed
    overhead for addresses, loop counters and operand fragments. Real numbers come from
    `ptxas -v` / Nsight; this is the first-order term that grows with the tile."""
    raise NotImplementedError("regs_per_thread")


def occupancy(gpu: GPU, cfg: TileConfig, in_bytes=2):
    """CTAs that can be resident on one SM, and which resource is the limit.

    Returns {"blocks": int, "limiter": str, "valid": bool, "reason": str}."""
    raise NotImplementedError("occupancy")


def waves(gpu: GPU, M, N, cfg: TileConfig, blocks_per_sm):
    """How many rounds of CTAs the grid needs, and how full the last round is (0..1]."""
    raise NotImplementedError("waves")


def roofline(gpu: GPU, flops, nbytes):
    """Lower-bound time and the bound that sets it."""
    raise NotImplementedError("roofline")


def analyze(gpu: GPU, M, N, K, cfg: TileConfig, in_bytes=2, out_bytes=2):
    """Everything above for one (shape, config) pair."""
    raise NotImplementedError("analyze")

import math

import pytest

from drill import load

k1 = load("k1_gemm_model")
G = k1.RTX3070
C = k1.TileConfig


def test_flops_and_min_bytes():
    assert k1.gemm_flops(4096, 4096, 4096) == 2 * 4096 ** 3
    # fp16 A, B, C: three 4096x4096 matrices x 2 bytes
    assert k1.gemm_min_bytes(4096, 4096, 4096) == 3 * 4096 * 4096 * 2
    assert k1.gemm_min_bytes(8, 16, 4, in_bytes=4, out_bytes=4) == (8 * 4 + 4 * 16 + 8 * 16) * 4


def test_tiled_bytes_rereads_operands():
    # 1024^3 with 128x128 tiles: A is read 1024/128 = 8 times, B 8 times
    got = k1.tiled_bytes(1024, 1024, 1024, C(128, 128, 64))
    assert got == (1024 * 1024 * 8 * 2) * 2 + 1024 * 1024 * 2
    # partial tiles count as a full re-read
    assert k1.tiled_bytes(100, 100, 10, C(64, 64, 16), in_bytes=1, out_bytes=1) == 100 * 10 * 2 * 2 + 100 * 100


def test_smem_regs_threads():
    cfg = C(128, 128, 64, num_warps=4, num_stages=3)
    assert k1.smem_bytes(cfg) == 3 * (128 * 64 + 64 * 128) * 2 == 98304
    assert k1.threads(cfg) == 128
    assert k1.regs_per_thread(cfg) == 128 * 128 // 128 + 40


def test_occupancy_limiters():
    big = k1.occupancy(G, C(128, 128, 64, 4, 3))
    assert big["valid"] and big["blocks"] == 1 and big["limiter"] == "smem"

    small = k1.occupancy(G, C(64, 64, 32, 4, 2))
    # threads 12, regs 65536 // (72*128) = 7, smem 102400 // 16384 = 6, blocks 16
    assert small["blocks"] == 6 and small["limiter"] == "smem"

    too_much_smem = k1.occupancy(G, C(128, 256, 64, 8, 3))
    assert not too_much_smem["valid"] and too_much_smem["limiter"] == "smem"

    spills = k1.occupancy(G, C(256, 256, 32, 4, 2))
    assert not spills["valid"] and spills["limiter"] == "regs"


def test_waves_and_tail():
    w = k1.waves(G, 4096, 4096, C(128, 128, 64), blocks_per_sm=1)
    assert w["tiles"] == 1024 and w["per_wave"] == 46 and w["waves"] == 23
    assert w["tail_fill"] == pytest.approx(12 / 46)
    exact = k1.waves(G, 46 * 64, 64, C(64, 64, 32), blocks_per_sm=1)
    assert exact["waves"] == 1 and exact["tail_fill"] == 1.0


def test_roofline_bounds():
    flops = k1.gemm_flops(4096, 4096, 4096)
    r = k1.roofline(G, flops, k1.gemm_min_bytes(4096, 4096, 4096))
    assert r["bound"] == "compute" and r["time"] == pytest.approx(flops / G.peak_flops)
    # a vector add: 1 flop per 6 bytes
    n = 1 << 24
    assert k1.roofline(G, n, 6 * n)["bound"] == "memory"
    # tiny problem: below the launch overhead
    tiny = k1.roofline(G, k1.gemm_flops(64, 64, 64), k1.gemm_min_bytes(64, 64, 64))
    assert tiny["bound"] == "launch" and tiny["time"] == G.launch_overhead


def test_analyze_bundle():
    a = k1.analyze(G, 2048, 2048, 2048, C(128, 128, 32, 4, 3))
    assert a["flops"] == 2 * 2048 ** 3
    assert a["intensity"] == pytest.approx(a["flops"] / a["tiled_bytes"])
    assert a["occupancy"]["valid"] and "waves" in a
    bad = k1.analyze(G, 2048, 2048, 2048, C(256, 256, 64, 4, 4))
    assert not bad["occupancy"]["valid"] and "waves" not in bad
    assert math.isclose(a["roofline"]["t_compute"], a["flops"] / G.peak_flops)

"""K2 · A benchmark harness you can trust — your implementation (run ./check.sh k2).

The same rules apply to a Triton kernel, a Luminal graph or a cuBLAS call:

  * warm up first (JIT compile, autotune, caches, clocks) and never time the warmup;
  * synchronize before reading the clock (GPU launches are asynchronous);
  * repeat, and report the median and the spread, not the best run;
  * keep compile / search time separate from steady-state run time;
  * only call a change faster when the difference is larger than the noise.

Everything takes an injectable `timer` and `sync`, so tests drive it with a fake clock and a
real run passes `time.perf_counter` and `torch.cuda.synchronize`.
"""
import statistics
import time


def percentile(xs, q):
    """Linear-interpolation percentile, q in [0, 100] (same as numpy's default)."""
    raise NotImplementedError("percentile")


def summarize(samples):
    """Robust statistics of a list of timings."""
    raise NotImplementedError("summarize")


def bench(fn, warmup=10, repeats=50, timer=time.perf_counter, sync=lambda: None):
    """Time fn() `repeats` times after `warmup` untimed calls.

    sync() runs before *starting* the clock (drain earlier async work) and before *stopping* it
    (wait for this call's work). Returns summarize(samples) plus the raw samples."""
    raise NotImplementedError("bench")


def compile_vs_run(make_fn, repeats=20, timer=time.perf_counter, sync=lambda: None):
    """Separate one-time cost (building / compiling / searching + the first call) from
    steady-state cost. make_fn() returns the callable; that construction is part of the
    one-time cost, exactly like `torch.compile`, a Triton JIT or a Luminal search."""
    raise NotImplementedError("compile_vs_run")


def outliers(samples, k=1.5):
    """Indices outside Tukey's fences [q1 - k*IQR, q3 + k*IQR]."""
    raise NotImplementedError("outliers")


def compare(base, new):
    """Is `new` faster than `base`? Both are summarize() results.

    speedup = base.median / new.median. The verdict is "faster" / "slower" only when the
    interquartile ranges do not overlap; otherwise "within noise"."""
    raise NotImplementedError("compare")


class FakeClock:
    """Deterministic clock for tests and for reasoning about the harness: each fn() call
    advances time by the next duration in `durations` (cycled)."""

    def __init__(self, durations):
        self.t = 0.0
        self.durations = list(durations)
        self.i = 0
        self.calls = 0

    def __call__(self):
        return self.t

    def work(self):
        self.t += self.durations[self.i % len(self.durations)]
        self.i += 1
        self.calls += 1

import pytest

from drill import load

k2 = load("k2_bench_harness")


def test_percentile_matches_numpy_linear():
    xs = [5, 1, 4, 2, 3]
    assert k2.percentile(xs, 0) == 1 and k2.percentile(xs, 100) == 5 and k2.percentile(xs, 50) == 3
    assert k2.percentile([1, 2, 3, 4], 25) == pytest.approx(1.75)


def test_summarize_basic():
    s = k2.summarize([1.0, 2.0, 3.0, 4.0, 100.0])
    assert s["median"] == 3.0 and s["min"] == 1.0 and s["max"] == 100.0
    assert s["q1"] == 2.0 and s["q3"] == 4.0 and s["iqr"] == 2.0
    assert s["mean"] == 22.0          # the mean is dragged by the outlier, the median is not


def test_bench_skips_warmup_and_syncs():
    # warmup calls take 10 s (compilation), steady calls take 1 s
    clock = k2.FakeClock([10.0] * 3 + [1.0] * 100)
    syncs = []
    r = k2.bench(clock.work, warmup=3, repeats=20, timer=clock, sync=lambda: syncs.append(clock.t))
    assert r["n"] == 20 and r["median"] == 1.0 and max(r["samples"]) == 1.0
    assert clock.calls == 23
    assert len(syncs) == 1 + 2 * 20     # once after warmup, before and after every timed call


def test_compile_vs_run():
    clock = k2.FakeClock([1.0])
    built = []

    def make():
        clock.t += 7.0                   # "compilation"
        built.append(1)
        return clock.work

    r = k2.compile_vs_run(make, repeats=5, timer=clock)
    assert built == [1]
    assert r["first_call"] == pytest.approx(8.0)
    assert r["steady_median"] == pytest.approx(1.0)
    assert r["one_time_cost"] == pytest.approx(7.0)


def test_outliers():
    xs = [1.0, 1.1, 0.9, 1.0, 1.05, 0.95, 5.0]
    assert k2.outliers(xs) == [6]


def test_compare_needs_more_than_noise():
    base = k2.summarize([10, 11, 12, 10.5, 11.5])
    clearly = k2.summarize([5, 5.5, 6, 5.2, 5.8])
    noisy = k2.summarize([9.5, 11, 12.5, 10, 11.8])
    assert k2.compare(base, clearly)["verdict"] == "faster"
    assert k2.compare(base, clearly)["speedup"] == pytest.approx(11 / 5.5)
    assert k2.compare(base, noisy)["verdict"] == "within noise"
    assert k2.compare(clearly, base)["verdict"] == "slower"

"""K3 · Numerical correctness: ULPs, tolerances, reassociation, edge shapes — your implementation (run ./check.sh k3).

Compilers reorder floating-point work all the time (split-K, tree reductions, fused epilogues,
e-graph rewrites like (a+b)+c -> a+(b+c)). The result is *different*, not necessarily *wrong*.
To tell the two apart you need: a distance in ULPs, a tolerance that scales with the dtype and the
reduction length, and inputs that hit the edges (non-divisible shapes, size 1, huge/tiny values).
"""
import math

import numpy as np

EPS = {"float32": 2.0 ** -23, "float16": 2.0 ** -10, "bfloat16": 2.0 ** -7}


def to_bf16(x):
    """Round float32 values to bfloat16 (round-to-nearest-even), returned as float32 values."""
    a = np.asarray(x, dtype=np.float32)
    bits = a.view(np.uint32).astype(np.uint64)
    rounding = ((bits >> 16) & 1) + 0x7FFF          # ties go to the even mantissa
    out = ((bits + rounding) >> 16) << 16
    return out.astype(np.uint32).view(np.float32)




def ulp_diff(a, b, dtype="float32"):
    """Number of representable `dtype` values between a and b (0 means identical)."""
    raise NotImplementedError("ulp_diff")


def sum_sequential(xs):
    raise NotImplementedError("sum_sequential")


def sum_pairwise(xs):
    """Tree reduction, the order a GPU block reduction (or tl.sum) uses."""
    raise NotImplementedError("sum_pairwise")


def tolerance(dtype, k=1, c=2.0):
    """(rtol, atol) for comparing a result computed in `dtype` with a high-precision reference.

    The error of a length-k reduction grows roughly like sqrt(k) * eps when the rounding errors
    are random, so rtol = c * eps * sqrt(k). atol covers values near zero, where relative error
    is meaningless; using the same number is a reasonable default for O(1) data."""
    raise NotImplementedError("tolerance")


def check_close(actual, expected, rtol, atol):
    """Like torch.testing.assert_close, but returns a report instead of raising.

    passes iff |actual - expected| <= atol + rtol * |expected| everywhere (NaNs must match)."""
    raise NotImplementedError("check_close")


def edge_sizes(block, limit=None):
    """Sizes that break tiled kernels: 1, block-1, block, block+1, 2*block+1, a prime, and
    (optionally) a large non-multiple. Sorted and de-duplicated."""
    raise NotImplementedError("edge_sizes")


def edge_values(dtype="float16"):
    """Values that stress a dtype: zeros of both signs, the smallest normal, the largest finite,
    values whose square overflows, and +/-inf / NaN."""
    raise NotImplementedError("edge_values")

import math

import numpy as np
import pytest

from drill import load

k3 = load("k3_numerics")


def test_bf16_rounding():
    x = np.array([1.0, 1.0 + 2 ** -8, 1.0 + 3 * 2 ** -8, 3.14159265], np.float32)
    got = k3.to_bf16(x)
    assert got[0] == 1.0
    assert got[1] == 1.0                      # exactly halfway -> even mantissa (1.0)
    assert got[2] == np.float32(1.0 + 2 ** -6)  # halfway between 1+2^-7 and 1+2^-6 -> even (1+2^-6)
    assert abs(got[3] - 3.140625) < 1e-7      # bf16 keeps 8 significant bits


def test_ulp_diff():
    one = np.float32(1.0)
    nxt = np.nextafter(one, np.float32(2.0))
    assert k3.ulp_diff(one, nxt) == 1
    assert k3.ulp_diff(one, one) == 0
    assert k3.ulp_diff(np.float32(-0.0), np.float32(0.0)) == 0
    tiny = np.nextafter(np.float32(0), np.float32(1))
    assert k3.ulp_diff(-tiny, tiny) == 2      # across zero
    assert k3.ulp_diff(np.float16(1.0), np.float16(1.0 + 2 ** -10), "float16") == 1
    assert k3.ulp_diff(1.0, 1.0 + 2 ** -7, "bfloat16") == 1


def test_reassociation_changes_the_bits():
    xs = np.array([1e8, 1.0, -1e8, 1.0] * 4, np.float32)
    seq, tree = k3.sum_sequential(xs), k3.sum_pairwise(xs)
    assert seq != tree                          # same math, different rounding
    assert k3.sum_pairwise([1.0, 2.0, 3.0]) == 6.0 and k3.sum_pairwise([]) == 0.0


def test_tolerance_scales():
    r1, a1 = k3.tolerance("float32", 1)
    r4, _ = k3.tolerance("float32", 4)
    assert r1 == a1 == 2 * 2 ** -23
    assert r4 == pytest.approx(2 * r1)
    assert k3.tolerance("bfloat16", 1)[0] > k3.tolerance("float16", 1)[0] > r1


def test_check_close_report():
    e = np.array([1.0, 2.0, 0.0, np.nan])
    ok = k3.check_close(e + 1e-9, e, rtol=1e-6, atol=1e-6)
    assert ok["ok"] and ok["mismatched"] == 0
    bad = k3.check_close(np.array([1.0, 2.1, 0.0, np.nan]), e, rtol=1e-3, atol=1e-6)
    assert not bad["ok"] and bad["mismatched"] == 1 and bad["index"] == (1,)
    nanmis = k3.check_close(np.array([1.0, 2.0, np.nan, np.nan]), e, 1e-3, 1e-3)
    assert not nanmis["ok"] and "NaN" in nanmis["reason"]
    assert not k3.check_close(np.zeros(3), np.zeros(4), 1, 1)["ok"]


def test_fp16_matmul_needs_a_tolerance():
    rng = np.random.default_rng(0)
    k = 512
    a, b = rng.standard_normal((64, k)), rng.standard_normal((k, 64))
    ref = a @ b
    low = (a.astype(np.float16) @ b.astype(np.float16)).astype(np.float64)
    assert not k3.check_close(low, ref, *k3.tolerance("float32", k))["ok"]
    rtol, atol = k3.tolerance("float16", k, c=4.0)
    assert k3.check_close(low, ref, rtol, atol * math.sqrt(k))["ok"]


def test_edge_sizes_and_values():
    s = k3.edge_sizes(64, limit=1000)
    assert s[0] == 1 and {63, 64, 65, 129, 97}.issubset(s) and s == sorted(set(s))
    assert any(x % 64 for x in s if x > 129)
    v = k3.edge_values("float16")
    assert np.isnan(v).any() and np.isinf(v).sum() == 2
    assert np.isinf(np.float16(v[5]) * np.float16(v[5]))    # its square overflows fp16

"""K4 · Fusion: fewer launches and less DRAM traffic vs more register pressure — your implementation (run ./check.sh k4).

A chain of ops is a list of stages. Unfused, every stage is its own kernel: it reads its inputs from
DRAM and writes its output back. Fused, a group of consecutive stages runs in one kernel: only the
group's external inputs are read and only its last output is written; the intermediates live in
registers. That is the Triton view (one kernel, an epilogue) and the Luminal view (a graph rewrite
that merges nodes) of the same trade-off.
"""
from dataclasses import dataclass


@dataclass(frozen=True)
class Stage:
    name: str
    out_elems: int          # elements written by this stage
    ext_in_elems: int       # elements read from outside the chain (weights, bias, residual, ...)
    flops: int
    live_regs: int = 0      # registers per thread this stage keeps live inside a fused kernel


def chain_mlp_epilogue(M, N, K):
    """GEMM followed by bias, GELU, residual add — the canonical fused epilogue."""
    return [
        Stage("matmul", M * N, M * K + K * N, 2 * M * N * K, live_regs=128),
        Stage("bias", M * N, N, M * N, live_regs=8),
        Stage("gelu", M * N, 0, 8 * M * N, live_regs=16),
        Stage("residual", M * N, M * N, M * N, live_regs=8),
    ]


def groups_of(stages, fuse):
    """fuse: list of group sizes summing to len(stages), e.g. [1, 1, 1, 1] (unfused) or [4]."""
    assert sum(fuse) == len(stages) and all(g > 0 for g in fuse)
    out, i = [], 0
    for g in fuse:
        out.append(stages[i:i + g])
        i += g
    return out


def traffic(stages, fuse, elem_bytes=2):
    """DRAM bytes, launches, and peak live registers for a fusion plan."""
    raise NotImplementedError("traffic")


def time_estimate(stages, fuse, peak_flops, peak_bw, launch=5e-6, elem_bytes=2, reg_limit=255, spill_penalty=2.0):
    """Roofline per group + launch overhead. If a fused group's live registers exceed reg_limit,
    its memory time is multiplied by spill_penalty (local-memory spills)."""
    raise NotImplementedError("time_estimate")


def best_plan(stages, peak_flops, peak_bw, **kw):
    """Try every way to cut the chain into consecutive groups (2^(n-1) plans)."""
    raise NotImplementedError("best_plan")

import pytest

from drill import load

k4 = load("k4_fusion")
S = 4096
PF, BW = 40e12, 448e9


def test_traffic_unfused_vs_fused():
    st = k4.chain_mlp_epilogue(S, S, S)
    mn = S * S
    un = k4.traffic(st, [1, 1, 1, 1])
    # matmul: A + B + C ; bias: C + b + C' ; gelu: 2 MN ; residual: 3 MN   (x 2 bytes)
    assert un["bytes"] == (2 * mn + mn + (mn + S + mn) + 2 * mn + 3 * mn) * 2
    assert un["launches"] == 4
    fused = k4.traffic(st, [4])
    assert fused["bytes"] == (2 * mn + S + mn + mn) * 2 and fused["launches"] == 1
    assert fused["peak_regs"] == 128 + 8 + 16 + 8
    assert k4.traffic(st, [1, 3])["bytes"] == k4.traffic(st, [2, 2])["bytes"]


def test_groups_validation():
    st = k4.chain_mlp_epilogue(8, 8, 8)
    with pytest.raises(AssertionError):
        k4.groups_of(st, [2, 1])


def test_best_plan_fuses_the_epilogue():
    t, plan = k4.best_plan(k4.chain_mlp_epilogue(S, S, S), PF, BW)
    assert plan == [4]
    t_un = k4.time_estimate(k4.chain_mlp_epilogue(S, S, S), [1, 1, 1, 1], PF, BW)
    assert t < t_un


def test_register_pressure_can_make_fusion_lose():
    n = 1 << 26
    heavy = [k4.Stage("a", n, n, n, 200), k4.Stage("b", n, 0, n, 100)]
    fused = k4.time_estimate(heavy, [2], PF, BW, spill_penalty=3.0)
    split = k4.time_estimate(heavy, [1, 1], PF, BW, spill_penalty=3.0)
    assert split < fused
    assert k4.best_plan(heavy, PF, BW, spill_penalty=3.0)[1] == [1, 1]
    light = [k4.Stage("a", n, n, n, 100), k4.Stage("b", n, 0, n, 100)]
    assert k4.best_plan(light, PF, BW, spill_penalty=3.0)[1] == [2]


def test_small_problem_is_launch_bound():
    st = k4.chain_mlp_epilogue(16, 64, 64)
    t_un = k4.time_estimate(st, [1, 1, 1, 1], PF, BW)
    t_f = k4.time_estimate(st, [4], PF, BW)
    assert t_un == pytest.approx(4 * 5e-6, rel=0.05) and t_f == pytest.approx(5e-6, rel=0.2)

"""T1 · IR dump lab: one kernel, four representations — your implementation (run ./check.sh t1).

Triton compiles a @triton.jit function through a fixed list of stages (see
third_party/nvidia/backend/compiler.py: make_ttir -> make_ttgir -> make_llir -> make_ptx -> make_cubin).
`kernel[grid](...)` returns a CompiledKernel whose `.asm` dict holds the text of each stage:

    k = add_kernel[grid](x, y, out, n, BLOCK=1024)
    k.asm["ttir"], k.asm["ttgir"], k.asm["llir"], k.asm["ptx"]

This drill is the *reading* half: given those texts (fixtures/t1/*), tell which stage a dump is,
count ops per dialect, pull out layout encodings and layout conversions, list the tensor-core
instructions in PTX, and diff which concepts appear / disappear between stages.
Runs anywhere — on a GPU box, `dump_stages()` produces the real texts.
"""
import re
from collections import Counter

# ---------------------------------------------------------------- stage detection
def detect_stage(text):
    """'ptx' | 'llir' | 'ttgir' | 'ttir' (checked in that order: lower stages are unambiguous)."""
    raise NotImplementedError("detect_stage")


# ---------------------------------------------------------------- MLIR ops
_OP = re.compile(r"^\s*(?:%[\w#:]+(?:\s*,\s*%[\w#:]+)*\s*=\s*)?\"?([a-z_]+\.[a-z_][\w.]*)\"?(?=[\s({<]|$)")


def mlir_ops(text):
    """Counter of fully qualified op names ('tt.load', 'ttg.convert_layout', 'scf.for', ...).
    One op per line in printed MLIR; attribute alias lines (#x = ...) and comments are skipped."""
    raise NotImplementedError("mlir_ops")


def dialect_histogram(text):
    """Counter dialect -> number of ops ('tt', 'ttg', 'arith', 'scf', ...)."""
    raise NotImplementedError("dialect_histogram")


# ---------------------------------------------------------------- layouts
def parse_attr_body(body):
    """'sizePerThread = [1, 4], order = [1, 0], vec = 8, parent = #blocked' -> dict.
    Lists become lists of ints, ints become ints, anything else stays a string."""
    raise NotImplementedError("parse_attr_body")


def layout_aliases(text):
    """{'#blocked': {'kind': 'blocked', 'sizePerThread': [...], ...}, '#mma': {'kind': 'nvidia_mma', ...}}"""
    raise NotImplementedError("layout_aliases")




def convert_layouts(text):
    """[(src_encoding, dst_encoding)] for every ttg.convert_layout."""
    raise NotImplementedError("convert_layouts")


# ---------------------------------------------------------------- PTX / LLVM
_PTX_KINDS = {
    "mma": r"\b((?:mma|wgmma\.mma_async)\.[\w.]+)",
    "ldmatrix": r"\b(ldmatrix\.[\w.]+)",
    "cp_async": r"\b(cp\.async\.[\w.]+|cp\.async\.bulk[\w.]*)",
    "barrier": r"\b(bar\.sync|barrier\.sync|mbarrier\.[\w.]+)",
    "global_vec": r"\b((?:ld|st)\.global(?:\.\w+)*\.v\d\.\w+)",
}


def ptx_instructions(text):
    """{'mma': Counter({'mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32': 2}), 'ldmatrix': ..., ...}
    Works on PTX and on LLVM IR (Triton emits these as inline asm strings there)."""
    raise NotImplementedError("ptx_instructions")


def mma_shape(instr):
    """'mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32' -> (16, 8, 16)."""
    raise NotImplementedError("mma_shape")


# ---------------------------------------------------------------- concepts per stage
CONCEPTS = {
    "program_id": r"tt\.get_program_id|ctaid",
    "block_tensor": r"tensor<",
    "pointer_tensor": r"tensor<[^>]*!tt\.ptr",
    "layout_encoding": r"#ttg\.(blocked|nvidia_mma|slice|dot_op)|, #\w+>",
    "shared_memory": r"ttg\.local_alloc|memdesc|addrspace\(3\)|\.shared",
    "async_copy": r"async_copy|cp\.async",
    "mma_instruction": r"mma\.sync|wgmma|ldmatrix",
    "thread_id": r"tid\.x|%tid",
    "virtual_registers": r"\.reg\s+\.",
    "high_level_dot": r"tt\.dot\b",
}


def concepts(text):
    raise NotImplementedError("concepts")


def concept_diff(stages):
    """stages: ordered list of (name, text). Returns [(from, to, appeared, disappeared)]."""
    raise NotImplementedError("concept_diff")


# ---------------------------------------------------------------- the GPU half
def dump_stages(compiled_kernel):
    """On a CUDA machine: compiled = kernel[grid](...); dump_stages(compiled) -> {stage: text}."""
    return {k: compiled_kernel.asm[k] for k in ("ttir", "ttgir", "llir", "ptx") if k in compiled_kernel.asm}

from pathlib import Path

from drill import load

t1 = load("t1_ir_anatomy")
FX = Path(__file__).resolve().parent.parent / "fixtures" / "t1"


def read(name):
    return (FX / name).read_text()


def test_detect_stage():
    assert t1.detect_stage(read("add.ttir.mlir")) == "ttir"
    assert t1.detect_stage(read("add.ttgir.mlir")) == "ttgir"
    assert t1.detect_stage(read("matmul.ttgir.mlir")) == "ttgir"
    assert t1.detect_stage(read("matmul.llir.ll")) == "llir"
    assert t1.detect_stage(read("matmul.ptx")) == "ptx"


def test_ops_and_dialects():
    ops = t1.mlir_ops(read("add.ttir.mlir"))
    assert ops["tt.load"] == 2 and ops["tt.store"] == 1 and ops["tt.splat"] == 5
    assert ops["tt.addptr"] == 3 and ops["arith.addf"] == 1 and ops["tt.return"] == 1
    assert ops["tt.func"] == 1
    d = t1.dialect_histogram(read("matmul.ttgir.mlir"))
    assert d["ttg"] >= 6 and d["scf"] == 2 and d["tt"] >= 3
    mm = t1.mlir_ops(read("matmul.ttgir.mlir"))
    assert mm["tt.dot"] == 1 and mm["ttg.local_load"] == 2 and mm["ttg.convert_layout"] == 1


def test_same_ops_ttir_vs_ttgir():
    # adding layouts must not change which ops exist in this simple kernel
    assert t1.mlir_ops(read("add.ttir.mlir")) == t1.mlir_ops(read("add.ttgir.mlir"))


def test_layout_aliases():
    a = t1.layout_aliases(read("matmul.ttgir.mlir"))
    assert a["#blocked"] == {"kind": "blocked", "sizePerThread": [1, 8], "threadsPerWarp": [4, 8],
                             "warpsPerCTA": [4, 1], "order": [1, 0]}
    assert a["#mma"]["kind"] == "nvidia_mma" and a["#mma"]["versionMajor"] == 2
    assert a["#mma"]["instrShape"] == [16, 8] and a["#mma"]["warpsPerCTA"] == [2, 2]
    assert a["#shared"]["kind"] == "swizzled_shared" and a["#shared"]["vec"] == 8
    assert a["#smem"] == {"kind": "shared_memory"}
    assert t1.layout_aliases(read("add.ttir.mlir")) == {}


def test_parse_attr_body_nested():
    d = t1.parse_attr_body("opIdx = 0, parent = #mma, kWidth = 2")
    assert d == {"opIdx": 0, "parent": "#mma", "kWidth": 2}


def test_convert_layouts():
    cv = t1.convert_layouts(read("matmul.ttgir.mlir"))
    assert cv == [("#mma", "#blocked1")]
    assert t1.convert_layouts(read("add.ttgir.mlir")) == []


def test_ptx_instructions():
    p = t1.ptx_instructions(read("matmul.ptx"))
    mma = p["mma"]
    assert sum(mma.values()) == 2
    (name,) = mma
    assert name == "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"
    assert t1.mma_shape(name) == (16, 8, 16)
    assert sum(p["ldmatrix"].values()) == 2 and "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16" in p["ldmatrix"]
    assert p["cp_async"]["cp.async.wait_group"] == 1 and p["barrier"]["bar.sync"] == 1
    assert p["global_vec"]["st.global.v4.b32"] == 1
    # the LLVM IR carries the same instructions as inline asm
    assert t1.ptx_instructions(read("matmul.llir.ll"))["mma"] == mma.__class__({name: 1})


def test_concept_diff_tells_the_story():
    stages = [(s, read(f)) for s, f in [("ttir", "matmul.ttir.mlir"), ("ttgir", "matmul.ttgir.mlir"),
                                        ("llir", "matmul.llir.ll"), ("ptx", "matmul.ptx")]]
    diff = {(a, b): (app, gone) for a, b, app, gone in t1.concept_diff(stages)}
    app, gone = diff[("ttir", "ttgir")]
    assert "layout_encoding" in app and "shared_memory" in app and "async_copy" in app and gone == []
    app, gone = diff[("ttgir", "llir")]
    assert "mma_instruction" in app and "thread_id" in app
    assert {"block_tensor", "layout_encoding", "high_level_dot", "program_id"} <= set(gone)
    app, gone = diff[("llir", "ptx")]
    assert "virtual_registers" in app

"""T2 · Layout reasoning: re-derive Triton's coalesce pass, then map lanes to elements — your implementation (run ./check.sh t2).

Triton's `tritongpu-coalesce` pass picks a BlockedEncoding for every load/store. Two pieces of
real code decide it, and this drill reimplements both in Python:

  1. lib/Dialect/TritonGPU/Transforms/Utility.cpp  getNumElementsPerThread()
        perThread = min(divisibility_bytes / elem_bytes, contiguity, shape[order[0]], 128 / elem_bits)
     then CoalesceUtils.cpp also caps it by  max(numElems / numThreads, 1).
  2. include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td  BlockedEncodingAttr builder:
        walk dims from the fastest (order[0]) and hand out lanes, then warps; the slowest dim
        takes whatever is left.

With the layout in hand we answer the T2 questions: which lane owns which coordinates, is the
first load of a warp coalesced, and does a ttg.convert_layout stay in registers, need warp
shuffles, or go through shared memory (lib/Analysis/Utility.cpp: cvtReordersRegisters /
cvtNeedsWarpShuffle / cvtNeedsSharedMemory).
"""
import re
from dataclasses import dataclass
from itertools import product
from math import prod


@dataclass(frozen=True)
class Blocked:
    size_per_thread: tuple
    threads_per_warp: tuple
    warps_per_cta: tuple
    order: tuple

    def __str__(self):
        f = lambda v: "[" + ", ".join(map(str, v)) + "]"
        return (f"#ttg.blocked<{{sizePerThread = {f(self.size_per_thread)}, threadsPerWarp = "
                f"{f(self.threads_per_warp)}, warpsPerCTA = {f(self.warps_per_cta)}, order = {f(self.order)}}}>")


def parse_blocked(text):
    """Parse '#ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], ...}>'."""
    def grab(key):
        m = re.search(key + r"\s*=\s*\[([^\]]*)\]", text)
        if not m:
            raise ValueError(f"missing {key}")
        return tuple(int(x) for x in m.group(1).split(","))
    return Blocked(grab("sizePerThread"), grab("threadsPerWarp"), grab("warpsPerCTA"), grab("order"))


def clamp(x, lo, hi):
    return max(lo, min(x, hi))


def build_blocked(shape, size_per_thread, order, num_warps, threads_per_warp=32):
    """Port of the BlockedEncodingAttr::get(shape, sizePerThread, order, numWarps, numThreadsPerWarp) builder."""
    raise NotImplementedError("build_blocked")


def elements_per_thread(shape, order, elem_bits, contiguity, divisibility_bytes, num_warps, threads_per_warp=32):
    """getNumElementsPerThread + the numElems/numThreads cap from CoalesceUtils.cpp."""
    raise NotImplementedError("elements_per_thread")


def coalesced_layout(shape, order, elem_bits, contiguity, divisibility_bytes=16, num_warps=4):
    """The layout the coalesce pass gives a load/store whose pointer has this contiguity along order[0]."""
    raise NotImplementedError("coalesced_layout")


# ---------------------------------------------------------------- who owns what
def _delinearize(idx, sizes, order):
    """Split a linear index into per-dim coordinates; order[0] varies fastest."""
    c = [0] * len(sizes)
    for d in order:
        c[d] = idx % sizes[d]
        idx //= sizes[d]
    return c


def tile(L: Blocked):
    return tuple(s * t * w for s, t, w in zip(L.size_per_thread, L.threads_per_warp, L.warps_per_cta))


def thread_elements(L: Blocked, shape, warp, lane):
    """Coordinates held by (warp, lane), in register order (inside a chunk first, then repetitions,
    both linearized along `order`)."""
    raise NotImplementedError("thread_elements")


def owners(L: Blocked, shape):
    """coord -> set of (warp, lane) holding it (more than one = broadcast / duplication)."""
    raise NotImplementedError("owners")


def lane_table(L: Blocked, shape, warp=0):
    """[(lane, [coords...])] for one warp — the T2 deliverable."""
    raise NotImplementedError("lane_table")


# ---------------------------------------------------------------- coalescing
def first_load_sectors(L: Blocked, shape, elem_bytes, strides=None):
    """Simulate warp 0's first vectorized load instruction on a tensor with element strides
    `strides` (row-major contiguous by default). Each lane loads one vector of
    vec = min(spt[order[0]], 16 // elem_bytes) contiguous elements along order[0].
    Returns (bytes_requested, sectors_touched, efficiency = bytes / (sectors * 32))."""
    raise NotImplementedError("first_load_sectors")


# ---------------------------------------------------------------- convert_layout cost
def classify_conversion(src: Blocked, dst: Blocked, shape):
    """'registers'      every element stays in the same thread(s): only registers are renamed
                         (Triton: cvtReordersRegisters)
       'warp_shuffle'   every element stays in the same warp: lanes exchange with shfl
                         (Triton: cvtNeedsWarpShuffle — the real check also limits the permutation)
       'shared_memory'  some element must cross warps: write to smem, barrier, read back
                         (Triton: cvtNeedsSharedMemory)"""
    raise NotImplementedError("classify_conversion")

from drill import load

t2 = load("t2_coalesce_layout")
B = t2.Blocked


def test_parse_roundtrip():
    s = "#ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [4, 1], order = [1, 0]}>"
    L = t2.parse_blocked(s)
    assert L == B((1, 4), (2, 16), (4, 1), (1, 0))
    assert str(L) == s


def test_reproduces_coalesce_mlir_transpose():
    # test/TritonGPU/coalesce.mlir @transpose: 64x64 f32, num-warps 4, tt.divisibility = 16
    row = t2.coalesced_layout([64, 64], order=[1, 0], elem_bits=32, contiguity=64)
    col = t2.coalesced_layout([64, 64], order=[0, 1], elem_bits=32, contiguity=64)
    assert str(row) == ("#ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], "
                        "warpsPerCTA = [4, 1], order = [1, 0]}>")
    assert str(col) == ("#ttg.blocked<{sizePerThread = [4, 1], threadsPerWarp = [16, 2], "
                        "warpsPerCTA = [1, 4], order = [0, 1]}>")


def test_elements_per_thread_limits():
    # fp16: 128 bits -> 8 elements per thread
    assert t2.elements_per_thread([128, 64], [1, 0], 16, 64, 16, 4) == 8
    # alignment 8 bytes only -> 2 fp32
    assert t2.elements_per_thread([128, 64], [1, 0], 32, 64, 8, 4) == 2
    # contiguity 1 (gather-like pointer) -> 1
    assert t2.elements_per_thread([128, 64], [1, 0], 32, 1, 16, 4) == 1
    # tiny tensor: 64 elements over 128 threads -> 1
    assert t2.elements_per_thread([64], [0], 32, 64, 16, 4) == 1


def test_build_blocked_1d_and_small():
    assert t2.build_blocked([1024], [4], [0], 4) == B((4,), (32,), (4,), (0,))
    # 16 columns of fp16 with 8 per thread: only 2 lanes along dim 1, the rest go to rows
    L = t2.build_blocked([64, 16], [1, 8], [1, 0], 4)
    assert L == B((1, 8), (16, 2), (4, 1), (1, 0))


def test_thread_elements_and_lane_table():
    L = B((1, 4), (2, 16), (4, 1), (1, 0))
    assert t2.tile(L) == (8, 64)
    assert t2.thread_elements(L, [8, 64], 0, 0) == [(0, 0), (0, 1), (0, 2), (0, 3)]
    assert t2.thread_elements(L, [8, 64], 0, 17) == [(1, 4), (1, 5), (1, 6), (1, 7)]
    assert t2.thread_elements(L, [8, 64], 3, 16)[0] == (7, 0)
    # 64x64: the tile repeats 8 times along rows -> 32 registers per thread
    regs = t2.thread_elements(L, [64, 64], 0, 0)
    assert len(regs) == 32 and regs[4] == (8, 0)
    table = t2.lane_table(L, [8, 64])
    assert len(table) == 32 and table[16][1][0] == (1, 0)


def test_broadcast_when_tensor_smaller_than_tile():
    L = B((1, 4), (2, 16), (4, 1), (1, 0))
    own = t2.owners(L, [4, 64])          # tile is 8 rows, tensor has 4
    assert len(own) == 4 * 64
    assert len(own[(0, 0)]) == 2          # warp 0 lane 0 and warp 2 lane 0
    assert own[(0, 0)] == {(0, 0), (2, 0)}


def test_coalescing():
    row = B((1, 4), (2, 16), (4, 1), (1, 0))
    n, sectors, eff = t2.first_load_sectors(row, [64, 64], 4)
    assert n == 32 * 16 and sectors == 16 and eff == 1.0
    # the same row layout reading a column-major (transposed) tensor
    n, sectors, eff = t2.first_load_sectors(row, [64, 64], 4, strides=[1, 64])
    assert sectors == 64 and eff == 0.25          # rows 0 and 1 share a sector, 64 columns do not
    # one element per lane, lanes along the slow dim -> 32 different rows
    bad = B((1, 1), (32, 1), (4, 1), (1, 0))
    assert t2.first_load_sectors(bad, [128, 64], 4)[1] == 32


def test_classify_conversion():
    shape = [64, 64]
    row = B((1, 4), (2, 16), (4, 1), (1, 0))
    assert t2.classify_conversion(row, row, shape) == "registers"
    # order also decides how lanes are numbered, so flipping it is NOT a register-only change
    flipped = B((1, 4), (2, 16), (4, 1), (0, 1))
    assert t2.classify_conversion(row, flipped, shape) != "registers"
    # each warp keeps the same rows; only which lane holds which column changes
    narrow = B((1, 2), (2, 16), (4, 1), (1, 0))
    assert t2.classify_conversion(row, narrow, shape) == "warp_shuffle"
    # lanes stacked 16 tall: warp 0 now owns rows 0-15, which warps 0-3 shared before
    tall = B((1, 4), (16, 2), (4, 1), (1, 0))
    assert t2.classify_conversion(row, tall, shape) == "shared_memory"
    col = B((4, 1), (16, 2), (1, 4), (0, 1))
    assert t2.classify_conversion(row, col, shape) == "shared_memory"

"""T3 · GEMM tile / resource sweep, the way triton.autotune does it — your implementation (run ./check.sh t3).

python/triton/runtime/autotuner.py runs, per new key:
    configs -> early_config_prune(configs, nargs)       # drop configs that cannot work
            -> perf_model + top_k                        # keep only the k most promising
            -> do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))  # measure the survivors
            -> pick the smallest median
We build the same pipeline in Python. The pruning and the perf model reuse K1's resource model
(solutions.k1_gemm_model, given), and `measure` is injectable: on the RTX 3070 it launches the
real Triton kernel, in tests it is a seeded fake.
"""
import csv
import io
import random
import statistics
from itertools import product
from math import ceil

from solutions.k1_gemm_model import RTX3070, TileConfig, occupancy, roofline, smem_bytes, tiled_bytes, gemm_flops


def config_space(BM=(32, 64, 128, 256), BN=(32, 64, 128, 256), BK=(32, 64), num_warps=(4, 8), num_stages=(2, 3, 4)):
    return [TileConfig(*c) for c in product(BM, BN, BK, num_warps, num_stages)]


def early_prune(configs, gpu, M, N, K):
    """Drop configs that cannot launch (K1 occupancy invalid) or whose tile is pointlessly larger
    than the problem (BM > 2 * M rounds a 100-row problem up to a 256-row tile)."""
    raise NotImplementedError("early_prune")


def perf_model(gpu, M, N, K, c, in_bytes=2):
    """Predicted seconds = max(compute with wave quantization, DRAM traffic) + launch overhead.

    Compute: tiles run in waves of sm_count * blocks_per_sm CTAs; a partly empty last wave
    still costs a full wave, so time = waves * per_wave * tile_flops / peak."""
    raise NotImplementedError("perf_model")


def top_k(configs, model, k):
    """The k configs with the smallest predicted time (stable on ties)."""
    raise NotImplementedError("top_k")


def quantiles(samples, qs=(0.5, 0.2, 0.8)):
    raise NotImplementedError("quantiles")


def sweep(configs, measure, repeats=25):
    """Measure each config `repeats` times. Rows: cfg fields + median/p20/p80/cv."""
    raise NotImplementedError("sweep")


def pick_best(rows):
    """Smallest median; if another row's interval [p20, p80] overlaps the winner's, prefer the
    one with less shared memory (more occupancy headroom, same speed within noise)."""
    raise NotImplementedError("pick_best")


def bound(gpu, M, N, K, c):
    """'memory' | 'compute' | 'launch' from the K1 roofline of this config's real traffic."""
    raise NotImplementedError("bound")


def to_csv(rows):
    buf = io.StringIO()
    w = csv.DictWriter(buf, fieldnames=list(rows[0]))
    w.writeheader()
    w.writerows(rows)
    return buf.getvalue()


def autotune(gpu, M, N, K, measure, configs=None, k=8, repeats=25):
    """The whole autotuner: prune -> perf_model top-k -> sweep -> pick. Returns (best_row, rows, n_pruned)."""
    raise NotImplementedError("autotune")


def fake_measure(gpu, M, N, K, seed=0, noise=0.03):
    """A seeded stand-in for a real launch: the perf model times a config-specific efficiency
    factor (what the model does not know) times multiplicative noise."""
    rng = random.Random(seed)

    def m(c):
        eff = 1.0 + 0.15 * ((hash((c.BM, c.BN, c.BK, c.num_warps, c.num_stages)) % 7) / 7)
        return perf_model(gpu, M, N, K, c) * eff * (1 + rng.gauss(0, noise))
    return m

import pytest

from drill import load
from solutions.k1_gemm_model import RTX3070 as G, TileConfig as C, occupancy

t3 = load("t3_autotune_sweep")
S = 4096


def test_space_and_prune():
    cs = t3.config_space()
    assert len(cs) == 4 * 4 * 2 * 2 * 3
    alive = t3.early_prune(cs, G, S, S, S)
    assert len(alive) == 146 and all(occupancy(G, c)["valid"] for c in alive)
    skinny = t3.early_prune(cs, G, 100, S, S)
    assert max(c.BM for c in skinny) == 128          # 256-row tiles are pruned for M = 100


def test_perf_model_wave_quantization():
    c = C(128, 128, 32, 4, 3)                          # 2 CTAs / SM (smem-limited) -> 92 per wave
    assert occupancy(G, c)["blocks"] == 2
    t = t3.perf_model(G, S, S, S, c)
    assert t == pytest.approx(0.004873388571428571)
    # the same tile on 92 tiles (exactly one wave) vs 93 tiles (a second, almost empty wave)
    one = t3.perf_model(G, 128 * 92, 128, 4096, c)
    two = t3.perf_model(G, 128 * 93, 128, 4096, c)
    assert two > one


def test_top_k_is_sorted_by_model():
    alive = t3.early_prune(t3.config_space(), G, S, S, S)
    m = lambda c: t3.perf_model(G, S, S, S, c)
    k = t3.top_k(alive, m, 5)
    assert len(k) == 5 and [m(c) for c in k] == sorted(m(c) for c in k)
    assert all(m(k[-1]) <= m(c) for c in alive if c not in k)


def test_quantiles_match_do_bench_defaults():
    med, p20, p80 = t3.quantiles(list(range(11)))
    assert (med, p20, p80) == (5, 2, 8)


def test_sweep_pick_and_csv():
    m = t3.fake_measure(G, S, S, S, seed=0)
    cands = [C(128, 256, 32, 8, 2), C(128, 128, 32, 4, 3), C(64, 64, 32, 4, 2)]
    rows = t3.sweep(cands, m, repeats=15)
    assert len(rows) == 3 and all(r["p20"] <= r["median"] <= r["p80"] for r in rows)
    best = t3.pick_best(rows)
    assert best["median"] <= min(r["median"] for r in rows) * 1.1
    text = t3.to_csv(rows)
    assert text.splitlines()[0].startswith("BM,BN,BK,num_warps,num_stages,median")
    assert len(text.strip().splitlines()) == 4


def test_pick_best_prefers_less_smem_within_noise():
    rows = [{"median": 1.00, "p20": 0.95, "p80": 1.05, "smem": 98304},
            {"median": 1.02, "p20": 0.97, "p80": 1.07, "smem": 32768},
            {"median": 1.50, "p20": 1.45, "p80": 1.55, "smem": 8192}]
    assert t3.pick_best(rows)["smem"] == 32768


def test_autotune_end_to_end():
    best, rows, pruned = t3.autotune(G, S, S, S, t3.fake_measure(G, S, S, S, seed=0), k=8)
    assert pruned == 46 and len(rows) == 8
    assert (best["BM"], best["BN"]) in {(128, 256), (256, 128)}


def test_bound():
    assert t3.bound(G, 64, 64, 64, C(32, 32, 32)) == "launch"
    # no L2 in the model: a 4096^3 fp16 GEMM with 128x128 tiles re-reads 2 GB -> "memory"
    assert t3.bound(G, S, S, S, C(128, 128, 32)) == "memory"

"""T4 · Correctness bug minimization: fuzz -> fail -> shrink -> regression test — your implementation (run ./check.sh t4).

The loop every compiler contributor runs before opening an issue:

  1. random differential test against a reference (fixed seed, edge-heavy shapes, several dtypes,
     non-contiguous strides);
  2. keep the first failing input;
  3. shrink it — smallest shape, simplest dtype, contiguous if possible — while it still fails;
  4. print it as a pytest regression (or a lit test once you know which pass is wrong).

`ddmin` is Zeller's delta debugging for lists (e.g. a list of IR lines or of ops in a graph);
`shrink_case` is the structured version for (shape, dtype, layout) cases.
"""
import random

import numpy as np

from solutions.k3_numerics import check_close, edge_sizes, to_bf16, tolerance


# ---------------------------------------------------------------- generic list shrinking
def ddmin(items, fails):
    """Smallest sub-list (1-minimal) for which fails(sub) is still True. fails(items) must be True."""
    raise NotImplementedError("ddmin")


def shrink_int(n, fails, lo=1):
    """Smallest m in [lo, n] with fails(m) (assumes fails(n)); halve first, then step down."""
    raise NotImplementedError("shrink_int")


# ---------------------------------------------------------------- structured cases
DTYPE_ORDER = ["float32", "float16", "bfloat16"]       # simplest first


def shrink_case(case, fails):
    """case = {'shape': [..], 'dtype': str, 'contiguous': bool}. Shrink until nothing changes."""
    raise NotImplementedError("shrink_case")


def random_case(rng, block=64):
    sizes = edge_sizes(block, limit=4 * block)
    return {"shape": [rng.choice(sizes), rng.choice(sizes)], "dtype": rng.choice(DTYPE_ORDER),
            "contiguous": rng.random() < 0.7}


def fuzz(check, n=200, seed=0, block=64):
    """Run `check(case) -> bool ok` on n random cases; return (first failing case or None, tried)."""
    raise NotImplementedError("fuzz")


# ---------------------------------------------------------------- the op under test
NP = {"float32": np.float32, "float16": np.float16, "bfloat16": np.float32}   # numpy has no bf16: values are rounded with to_bf16


def make_input(case, seed=0):
    m, n = case["shape"]
    rng = np.random.default_rng(seed)
    if case["contiguous"]:
        x = rng.standard_normal((m, n))
    else:
        x = rng.standard_normal((n, m)).T            # same values pattern, transposed strides
    x = x.astype(NP[case["dtype"]])
    return to_bf16(x) if case["dtype"] == "bfloat16" else x


def row_sum_reference(x):
    return np.asarray(x, np.float64).sum(axis=1)


def differential_check(impl, case, seed=0):
    """impl(x) vs float64 reference with a K3 tolerance for this dtype and row length."""
    raise NotImplementedError("differential_check")


def buggy_row_sum(x, BLOCK=64):
    """A 'tiled kernel' with a classic bug: the single-block path is masked, the multi-block path
    walks only full blocks and forgets the tail."""
    m, n = x.shape
    acc = np.zeros(m, np.float32)
    if n <= BLOCK:
        acc += x.astype(np.float32).sum(axis=1)
    else:
        for k0 in range(0, n - BLOCK + 1, BLOCK):     # BUG: drops n % BLOCK trailing columns
            acc += x[:, k0:k0 + BLOCK].astype(np.float32).sum(axis=1)
    return acc


def to_pytest(case, name="test_row_sum_regression", impl="row_sum", seed=0):
    """Emit a standalone regression test for the minimized case."""
    raise NotImplementedError("to_pytest")

import numpy as np

from drill import load

t4 = load("t4_minimize")


def test_ddmin_finds_the_two_culprits():
    items = list(range(20))
    fails = lambda xs: 3 in xs and 17 in xs
    assert sorted(t4.ddmin(items, fails)) == [3, 17]
    assert t4.ddmin(list("abcdefgh"), lambda xs: "e" in xs) == ["e"]


def test_ddmin_ir_lines():
    ir = ["%0 = tt.get_program_id x", "%1 = tt.make_range", "%2 = tt.load %p", "%3 = tt.dot %a, %b",
          "%4 = ttg.convert_layout %3", "tt.store %q, %4", "tt.return"]
    crash = lambda lines: any("tt.dot" in l for l in lines) and any("convert_layout" in l for l in lines)
    assert t4.ddmin(ir, crash) == ["%3 = tt.dot %a, %b", "%4 = ttg.convert_layout %3"]


def test_shrink_int():
    assert t4.shrink_int(1000, lambda m: m >= 65) == 65
    assert t4.shrink_int(1000, lambda m: m % 64 != 0 and m > 64) == 65
    assert t4.shrink_int(7, lambda m: True) == 1


def test_reference_impl_passes_and_bug_is_found():
    good = lambda x: x.astype(np.float32).sum(axis=1)
    for case in [{"shape": [3, 129], "dtype": "float16", "contiguous": False},
                 {"shape": [1, 1], "dtype": "float32", "contiguous": True}]:
        assert t4.differential_check(good, case)
    bad, tried = t4.fuzz(lambda c: t4.differential_check(t4.buggy_row_sum, c), n=200, seed=0)
    assert bad is not None and bad["shape"][1] > 64 and bad["shape"][1] % 64 != 0


def test_shrink_case_to_minimal():
    check = lambda c: t4.differential_check(t4.buggy_row_sum, c)
    start = {"shape": [129, 255], "dtype": "bfloat16", "contiguous": False}
    assert not check(start)
    m = t4.shrink_case(start, lambda c: not check(c))
    assert m == {"shape": [1, 65], "dtype": "float32", "contiguous": True}


def test_to_pytest_runs():
    code = t4.to_pytest({"shape": [1, 65], "dtype": "float32", "contiguous": True}, impl="buggy")
    ns = {"np": np, "buggy": t4.buggy_row_sum}
    exec(code, ns)
    try:
        ns["test_row_sum_regression"]()
    except AssertionError:
        pass
    else:
        raise AssertionError("the emitted regression test should fail on the buggy kernel")
    ns["buggy"] = lambda x: x.astype(np.float32).sum(axis=1)
    exec(code, ns)
    ns["test_row_sum_regression"]()

"""T5 · A mini FileCheck: how Triton's lit tests decide pass / fail — your implementation (run ./check.sh t5).

Every file under triton/test/ starts with a RUN line such as
    // RUN: triton-opt %s -split-input-file -tritongpu-coalesce | FileCheck %s
lit runs the pass, pipes the output to LLVM's FileCheck, and FileCheck matches the `// CHECK...`
comments of the same file against that output. Supported here (the subset Triton tests use most):

  CHECK:        match somewhere at or after the current position
  CHECK-NEXT:   match on the line right after the previous match
  CHECK-SAME:   match on the same line as the previous match, after it
  CHECK-NOT:    must not appear between the previous match and the next positive match
  CHECK-LABEL:  splits input into blocks first; checks never cross a label
  CHECK-DAG:    a group of consecutive DAG checks may match in any order
  CHECK-COUNT-n: match n times in a row (each after the previous)

Patterns: literal text, except {{regex}} and [[VAR:regex]] (capture) / [[VAR]] (use), with
the `$` prefix ([[$VAR]]) meaning "global": with --enable-var-scope only those survive a CHECK-LABEL. Whitespace runs in the pattern
match any whitespace run in the input (FileCheck's default canonicalization).
"""
import re

DIRECTIVE = re.compile(r"(?://|;|#)\s*(CHECK(?:-NEXT|-SAME|-NOT|-LABEL|-DAG|-COUNT-\d+)?):(.*)$")


class CheckFailure(AssertionError):
    pass


def parse_checks(source, prefix="CHECK"):
    """[(kind, pattern, line_no)] from a test file. kind: 'CHECK', 'NEXT', 'SAME', 'NOT', 'LABEL', 'DAG', 'COUNT'."""
    raise NotImplementedError("parse_checks")


def compile_pattern(pat, env):
    """Turn a FileCheck pattern into a Python regex. Returns (regex, [capture variable names])."""
    raise NotImplementedError("compile_pattern")










def filecheck(check_source, output, prefix="CHECK", enable_var_scope=False):
    """Raise CheckFailure (an AssertionError) if `output` does not satisfy the CHECKs in check_source."""
    raise NotImplementedError("filecheck")


def split_input_file(text):
    """triton-opt -split-input-file: chunks separated by a line `// -----`."""
    return [c.strip("\n") for c in re.split(r"^// -----\s*$", text, flags=re.M)]

import pytest

from drill import load

t5 = load("t5_filecheck")
FC = t5.filecheck
Fail = t5.CheckFailure

OUT = """#blocked = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [4, 1], order = [1, 0]}>
#blocked1 = #ttg.blocked<{sizePerThread = [4, 1], threadsPerWarp = [16, 2], warpsPerCTA = [1, 4], order = [0, 1]}>
module {
  tt.func @transpose(%arg0: !tt.ptr<f32>) {
    %8 = ttg.convert_layout %7 : tensor<64x64x!tt.ptr<f32>, #blocked2> -> tensor<64x64x!tt.ptr<f32>, #blocked>
    %9 = tt.load %8 : tensor<64x64x!tt.ptr<f32>, #blocked>
    %10 = ttg.convert_layout %9 : tensor<64x64xf32, #blocked> -> tensor<64x64xf32, #blocked1>
    tt.store %11, %10 : tensor<64x64x!tt.ptr<f32>, #blocked1>
    tt.return
  }
  tt.func @other() {
    %0 = arith.constant 0 : i32
    tt.return
  }
}"""


def test_parse_checks():
    src = "// CHECK: a\n// CHECK-NEXT: b\n  // CHECK-COUNT-2: c\n// CHECK-NOT: d\n// CHECK-LABEL: @f\n; CHECK-SAME: e"
    kinds = [k for k, _, _ in t5.parse_checks(src)]
    assert kinds == ["CHECK", "NEXT", "CHECK", "CHECK", "NOT", "LABEL", "SAME"]
    assert t5.parse_checks("// MMA: x", prefix="MMA") == [("CHECK", "x", 1)]


def test_compile_pattern():
    rx, names = t5.compile_pattern("[[L:#.*]] = #ttg.blocked<{{.*}}order = [1, 0]", {})
    assert names == ["L"] and rx.search(OUT).group("L") == "#blocked"
    rx, _ = t5.compile_pattern("tt.load   [[P]]", {"P": "%8"})
    assert rx.search("%9 = tt.load %8 : x")


def test_real_coalesce_style_checks_pass():
    checks = """
// CHECK: [[row:#.*]] = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16]
// CHECK: [[col:#.*]] = #ttg.blocked<{sizePerThread = [4, 1]
// CHECK-LABEL: @transpose
// CHECK: [[ptr:%.*]] = ttg.convert_layout {{.*}} -> tensor<64x64x!tt.ptr<f32>, [[row]]>
// CHECK-NEXT: [[val:%[0-9]+]] = tt.load [[ptr]]
// CHECK-SAME: [[row]]>
// CHECK-NOT: tt.dot
// CHECK: tt.store {{.*}}, [[col]]>
"""
    assert FC(checks, OUT)


def test_failures_are_reported():
    with pytest.raises(Fail, match="not found"):
        FC("// CHECK: tt.dot", OUT)
    with pytest.raises(Fail, match="CHECK-NEXT"):
        FC("// CHECK: tt.func @transpose\n// CHECK-NEXT: tt.load", OUT)
    with pytest.raises(Fail, match="CHECK-NOT"):
        FC("// CHECK: @transpose\n// CHECK-NOT: tt.load\n// CHECK: tt.store", OUT)
    with pytest.raises(Fail, match="no check strings"):
        FC("// nothing here", OUT)


def test_order_matters_unless_dag():
    with pytest.raises(Fail):
        FC("// CHECK: tt.store\n// CHECK: tt.load", OUT)
    assert FC("// CHECK-DAG: tt.store\n// CHECK-DAG: tt.load", OUT)


def test_labels_bound_the_search():
    # arith.constant exists, but only in @other: it must not satisfy a check in the @transpose block
    with pytest.raises(Fail):
        FC("// CHECK-LABEL: @transpose\n// CHECK: arith.constant\n// CHECK-LABEL: @other", OUT)
    assert FC("// CHECK-LABEL: @transpose\n// CHECK: tt.return\n// CHECK-LABEL: @other\n// CHECK: arith.constant", OUT)


def test_count_and_variables():
    assert FC("// CHECK-COUNT-2: ttg.convert_layout", OUT)
    with pytest.raises(Fail):
        FC("// CHECK-COUNT-3: ttg.convert_layout", OUT)
    with pytest.raises(Fail, match="undefined"):
        FC("// CHECK: tt.load [[NOPE]]", OUT)


def test_var_scope():
    src = "// CHECK: [[B:#blocked1]] =\n// CHECK-LABEL: @transpose\n// CHECK: tt.store {{.*}}[[B]]>"
    assert FC(src, OUT)
    with pytest.raises(Fail, match="undefined"):
        FC(src, OUT, enable_var_scope=True)
    assert FC(src.replace("[[B", "[[$B"), OUT, enable_var_scope=True)


def test_split_input_file():
    parts = t5.split_input_file("a\n// -----\nb\nc\n// -----\nd")
    assert parts == ["a", "b\nc", "d"]

"""L1 · Graph semantics: y = relu(x @ w + b) as a Luminal-style primitive graph — your implementation (run ./check.sh l1).

Luminal records every tensor operation into a static DAG instead of running it (lazy, ahead of
time). High-level ops are *spelled* with a small primitive set (README: "Everything in Luminal
boils down to 15 primitive ops"; src/logical/ir.rs `LogicalOp` has the current full list):

  matmul(a, b)  = broadcast both to [M, N, K]  ->  Mul  ->  ReduceSum(K)     (src/frontend/matmul.rs)
  relu(x)       = maximum(x, 0) = Select(LessThan(x, 0), 0, x)                (src/frontend/unary.rs, binary.rs)
  exp(x)        = Exp2(x * log2(e))                                           (src/frontend/unary.rs)
  a + b[N]      = Add(a, Expand(b))  — broadcasting is a *view*, not a copy

This file is a ~150-line Python model of that design: a Graph of nodes with shapes and dtypes,
a GraphTensor API that records, a node-semantics table, and a reference runtime (numpy) that
executes the DAG in topological order — the analogue of Luminal's ReferenceRuntime.
"""
from dataclasses import dataclass, field

import numpy as np

PRIMITIVES = {"Input", "Constant", "Add", "Mul", "Recip", "Exp2", "Log2", "Sqrt", "Sin", "LessThan",
              "Select", "ReduceSum", "ReduceMax", "Expand", "Permute", "Cast"}
MOVEMENT = {"Expand", "Permute"}          # views: change indexing, move no data


@dataclass
class Node:
    id: int
    op: str
    inputs: list
    shape: tuple
    dtype: str = "f32"
    attrs: dict = field(default_factory=dict)


class Graph:
    def __init__(self):
        self.nodes = []
        self.outputs = []

    def add(self, op, inputs, shape, dtype="f32", **attrs):
        assert op in PRIMITIVES, op
        n = Node(len(self.nodes), op, [t.id for t in inputs], tuple(shape), dtype, attrs)
        self.nodes.append(n)
        return Tensor(self, n.id)

    def tensor(self, shape, name, dtype="f32"):
        return self.add("Input", [], shape, dtype, name=name)

    def constant(self, value, dtype="f32"):
        return self.add("Constant", [], (), dtype, value=float(value))

    # ---- inspection
    def topo_order(self):
        """Node ids such that every input comes before its user (ids already are; we verify)."""
        raise NotImplementedError("Graph.topo_order")

    def op_counts(self):
        raise NotImplementedError("Graph.op_counts")

    def semantics_table(self):
        """[(id, op, input ids, shape, dtype, note)] — the L1 deliverable."""
        raise NotImplementedError("Graph.semantics_table")

    def materialized_bytes(self, elem_bytes=4):
        """Bytes a naive executor would allocate: every non-view, non-input node gets a buffer."""
        raise NotImplementedError("Graph.materialized_bytes")


@dataclass(frozen=True)
class Tensor:
    g: Graph
    id: int

    @property
    def shape(self):
        return self.g.nodes[self.id].shape

    @property
    def dtype(self):
        return self.g.nodes[self.id].dtype

    def output(self):
        self.g.outputs.append(self.id)
        return self

    # ---- movement
    def expand(self, axis, size):
        s = list(self.shape)
        s.insert(axis, size)
        return self.g.add("Expand", [self], s, self.dtype, axis=axis, size=size)

    def expand_to(self, shape):
        """Right-aligned numpy-style broadcast by inserting leading / size-1-free axes."""
        t = self
        for axis in range(len(shape) - len(self.shape)):
            t = t.expand(axis, shape[axis])
        assert t.shape == tuple(shape), (t.shape, shape)
        return t

    def permute(self, axes):
        return self.g.add("Permute", [self], [self.shape[a] for a in axes], self.dtype, axes=tuple(axes))

    # ---- binary (shapes must already agree, like Luminal: broadcasting is explicit)
    def _bin(self, op, other, dtype=None):
        if not isinstance(other, Tensor):
            other = self.g.constant(other).expand_to(self.shape) if self.shape else self.g.constant(other)
        assert self.shape == other.shape, f"{op}: {self.shape} vs {other.shape}"
        return self.g.add(op, [self, other], self.shape, dtype or self.dtype)

    def __add__(self, o):
        return self._bin("Add", o)

    def __mul__(self, o):
        return self._bin("Mul", o)

    def lt(self, o):
        return self._bin("LessThan", o, "bool")

    def select(self, a, b):
        """self is the boolean condition."""
        assert self.shape == a.shape == b.shape
        return self.g.add("Select", [self, a, b], a.shape, a.dtype)

    def sum(self, axis):
        s = list(self.shape)
        s.pop(axis)
        return self.g.add("ReduceSum", [self], s, self.dtype, axis=axis)

    # ---- high-level ops, spelled with primitives
    def matmul(self, w):
        """[M, K] @ [K, N]: both operands become [M, N, K] views, then Mul and ReduceSum(K)."""
        raise NotImplementedError("Tensor.matmul")

    def relu(self):
        raise NotImplementedError("Tensor.relu")

    def exp(self):
        raise NotImplementedError("Tensor.exp")

    def exp2(self):
        return self.g.add("Exp2", [self], self.shape, self.dtype)


def build_mlp(M, K, N):
    """y = relu(x @ w + b). Returns (graph, x, w, b, y)."""
    raise NotImplementedError("build_mlp")


# ---------------------------------------------------------------- reference runtime
def run(g, feeds):
    """Execute in topological order with numpy. feeds: {input name: array}. Returns {node id: array}."""
    raise NotImplementedError("run")

import numpy as np
import pytest

from drill import load

l1 = load("l1_graph")


def test_recording_does_not_compute():
    g = l1.Graph()
    x = g.tensor((2, 3), "x")
    y = x + x
    assert isinstance(y, l1.Tensor) and len(g.nodes) == 2 and g.nodes[1].op == "Add"


def test_matmul_is_broadcast_mul_reduce():
    g = l1.Graph()
    a, b = g.tensor((5, 3), "a"), g.tensor((3, 7), "b")
    c = a.matmul(b).output()
    assert c.shape == (5, 7)
    ops = [g.nodes[i].op for i in g.topo_order()]
    assert ops.count("Mul") == 1 and ops.count("ReduceSum") == 1 and ops.count("Permute") == 1
    mul = next(n for n in g.nodes if n.op == "Mul")
    assert mul.shape == (5, 7, 3)                     # the [M, N, K] product
    with pytest.raises(AssertionError):
        a.matmul(g.tensor((4, 7), "bad"))


def test_mlp_graph_shape_and_counts():
    g, x, w, b, y = l1.build_mlp(2, 3, 4)
    assert y.shape == (2, 4) and y.dtype == "f32"
    assert g.op_counts() == {"Input": 3, "Expand": 5, "Permute": 1, "Mul": 1, "ReduceSum": 1, "Add": 1,
                             "Constant": 1, "LessThan": 1, "Select": 1}
    table = g.semantics_table()
    lt = next(r for r in table if r[1] == "LessThan")
    assert lt[4] == "bool"
    assert any("view" in r[5] for r in table if r[1] == "Expand")
    assert all(set(g.nodes[r[0]].op for r in table) <= l1.PRIMITIVES for _ in [0])


def test_reference_runtime_matches_numpy():
    rng = np.random.default_rng(0)
    for M, K, N in [(2, 3, 4), (1, 1, 1), (7, 13, 5)]:
        g, x, w, b, y = l1.build_mlp(M, K, N)
        X, W, B = rng.standard_normal((M, K)), rng.standard_normal((K, N)), rng.standard_normal(N)
        v = l1.run(g, {"x": X, "w": W, "b": B})
        np.testing.assert_allclose(v[y.id], np.maximum(X @ W + B, 0), rtol=1e-5, atol=1e-5)


def test_relu_is_select_and_exp_is_exp2():
    g = l1.Graph()
    x = g.tensor((3,), "x")
    r = x.relu().output()
    assert g.nodes[r.id].op == "Select"
    e = x.exp().output()
    assert g.nodes[e.id].op == "Exp2"
    v = l1.run(g, {"x": np.array([-1.0, 0.0, 2.0])})
    np.testing.assert_allclose(v[r.id], [0, 0, 2])
    np.testing.assert_allclose(v[e.id], np.exp([-1.0, 0.0, 2.0]), rtol=1e-6)


def test_naive_execution_materializes_the_product():
    g, *_ = l1.build_mlp(64, 256, 64)
    # Mul [64, 64, 256] alone is 4 MiB of fp32 — that's what fusion must avoid
    assert g.materialized_bytes() >= 64 * 64 * 256 * 4
    assert g.materialized_bytes() == (64 * 64 * 256 + 4 * 64 * 64) * 4

"""L2 · E-graphs and equality saturation from scratch — your implementation (run ./check.sh l2).

Luminal's search space is an e-graph built by egglog (src/egglog_core/egglog_preamble.egg holds
the rules). An e-graph stores *many equivalent programs at once*:

  e-node   an operator applied to e-classes      ("+", c1, c2)
  e-class  a set of e-nodes known to be equal    {a+b, b+a}
  rewrite  lhs -> rhs: when lhs matches in class C, add rhs and union it into C (nothing is deleted)
  rebuild  restore congruence: if a ≡ b then f(a) ≡ f(b)
  saturate apply all rules until nothing new appears (or a node / iteration budget runs out)
  extract  pick one e-node per class to minimize a cost -> one concrete program

Terms are tuples: ("+", x, y), leaves are strings or numbers. Pattern variables start with "?".
"""
from itertools import count


class EGraph:
    def __init__(self):
        self.parent = {}            # union-find over class ids
        self.classes = {}           # canonical id -> set of canonical e-nodes
        self.hashcons = {}          # canonical e-node -> class id
        self._ids = count()
        self.pending = []

    # ---------------------------------------------------------- union-find
    def find(self, a):
        while self.parent[a] != a:
            self.parent[a] = self.parent[self.parent[a]]
            a = self.parent[a]
        return a

    def canon(self, enode):
        op, *kids = enode
        return (op, *(self.find(k) for k in kids))

    # ---------------------------------------------------------- building
    def add_enode(self, enode):
        enode = self.canon(enode)
        if enode in self.hashcons:
            return self.find(self.hashcons[enode])
        cid = next(self._ids)
        self.parent[cid] = cid
        self.classes[cid] = {enode}
        self.hashcons[enode] = cid
        return cid

    def add(self, term):
        if isinstance(term, tuple):
            return self.add_enode((term[0], *(self.add(t) for t in term[1:])))
        return self.add_enode((term,))

    def union(self, a, b):
        raise NotImplementedError("EGraph.union")

    def rebuild(self):
        """Re-canonicalize every e-node; merge classes whose e-nodes became identical (congruence)."""
        raise NotImplementedError("EGraph.rebuild")

    # ---------------------------------------------------------- queries
    def equiv(self, t1, t2):
        raise NotImplementedError("EGraph.equiv")

    def num_classes(self):
        return len(self.classes)

    def num_nodes(self):
        return sum(len(v) for v in self.classes.values())

    # ---------------------------------------------------------- e-matching
    def ematch(self, pattern, cid, subst=None):
        """All substitutions {?var: class id} under which `pattern` matches something in class cid."""
        raise NotImplementedError("EGraph.ematch")

    def instantiate(self, pattern, subst):
        raise NotImplementedError("EGraph.instantiate")

    # ---------------------------------------------------------- saturation
    def run(self, rules, iter_limit=20, node_limit=10_000):
        """Equality saturation. rules: [(name, lhs, rhs)]. Returns a report dict."""
        raise NotImplementedError("EGraph.run")

    # ---------------------------------------------------------- extraction
    def extract(self, cid, cost):
        """Cheapest term in class cid. cost(op, child_costs) -> number. Bottom-up fixpoint
        (classes may be cyclic, e.g. x*1 = x)."""
        raise NotImplementedError("EGraph.extract")


# ---------------------------------------------------------------- rule sets
COMM_ASSOC = [
    ("comm-add", ("+", "?a", "?b"), ("+", "?b", "?a")),
    ("assoc-add", ("+", ("+", "?a", "?b"), "?c"), ("+", "?a", ("+", "?b", "?c"))),
]
ALGEBRA = COMM_ASSOC + [
    ("mul-one", ("*", "?a", 1), "?a"),
    ("add-zero", ("+", "?a", 0), "?a"),
    ("factor", ("+", ("*", "?a", "?b"), ("*", "?a", "?c")), ("*", "?a", ("+", "?b", "?c"))),
    ("exp2-log2", ("exp2", ("log2", "?a")), "?a"),
]


def ast_size(op, kids):
    return 1 + sum(kids)


def op_cost(op, kids):
    """Rough kernel cost: multiplies cost more than adds; leaves are free."""
    return {"*": 4, "+": 1, "exp2": 8, "log2": 8}.get(op, 0) + sum(kids)


def evaluate(term, env, dtype=float):
    """Evaluate a term with Python floats or numpy scalars — to show which rewrites change bits."""
    import math
    if isinstance(term, tuple):
        op, *k = term
        v = [evaluate(t, env, dtype) for t in k]
        return {"+": lambda a, b: dtype(a + b), "*": lambda a, b: dtype(a * b),
                "exp2": lambda a: dtype(2.0 ** a), "log2": lambda a: dtype(math.log2(a))}[op](*v)
    return dtype(env[term]) if isinstance(term, str) else dtype(term)

import numpy as np

from drill import load

l2 = load("l2_egraph")
E = l2.EGraph


def test_hashcons_shares_subterms():
    g = E()
    a = g.add(("+", "a", "b"))
    b = g.add(("+", "a", "b"))
    assert a == b and g.num_classes() == 3 and g.num_nodes() == 3


def test_union_and_congruence():
    g = E()
    fa, fb = g.add(("f", "a")), g.add(("f", "b"))
    assert g.find(fa) != g.find(fb)
    g.union(g.add("a"), g.add("b"))
    g.rebuild()
    assert g.find(fa) == g.find(fb)          # a = b  =>  f(a) = f(b)


def test_ematch():
    g = E()
    root = g.add(("+", ("+", "a", "b"), "c"))
    subs = g.ematch(("+", ("+", "?x", "?y"), "?z"), root)
    assert len(subs) == 1 and g.find(subs[0]["?x"]) == g.find(g.add("a"))
    assert g.ematch(("+", "?x", "?x"), root) == []
    assert len(g.ematch(("+", "?x", "?x"), g.add(("+", "a", "a")))) == 1


def test_three_term_saturation_counts():
    g = E()
    g.add(("+", ("+", "a", "b"), "c"))
    rep = g.run(l2.COMM_ASSOC)
    assert rep["stop"] == "saturated"
    # classes: a, b, c, a+b, a+c, b+c, a+b+c ; e-nodes: 3 leaves + 2 per pair + 6 at the top
    assert g.num_classes() == 7 and g.num_nodes() == 15
    assert g.equiv(("+", "a", ("+", "b", "c")), ("+", ("+", "a", "b"), "c"))
    assert g.equiv(("+", ("+", "c", "a"), "b"), ("+", "b", ("+", "c", "a")))
    assert not g.equiv(("+", "a", "b"), ("+", "a", "c"))


def test_blowup_is_exponential():
    sizes = []
    for n in (3, 4, 5, 6):
        g = E()
        t = "x0"
        for i in range(1, n):
            t = ("+", t, f"x{i}")
        g.add(t)
        assert g.run(l2.COMM_ASSOC, iter_limit=30, node_limit=10 ** 6)["stop"] == "saturated"
        assert g.num_classes() == 2 ** n - 1
        assert g.num_nodes() == 3 ** n - 2 ** (n + 1) + 1 + n
        sizes.append(g.num_nodes())
    assert sizes == [15, 54, 185, 608]


def test_budgets_stop_the_search():
    g = E()
    t = "x0"
    for i in range(1, 8):
        t = ("+", t, f"x{i}")
    g.add(t)
    assert g.run(l2.COMM_ASSOC, iter_limit=2)["stop"] == "iter_limit"
    g2 = E()
    g2.add(t)
    assert g2.run(l2.COMM_ASSOC, iter_limit=50, node_limit=500)["stop"] == "node_limit"


def test_extraction_picks_the_cheaper_program():
    g = E()
    root = g.add(("+", ("*", "x", "y"), ("*", "x", "z")))
    g.run(l2.ALGEBRA)
    cost, term = g.extract(root, l2.op_cost)
    assert cost == 5 and term[0] == "*" and g.equiv(term, ("*", "x", ("+", "y", "z")))
    g = E()
    root = g.add(("*", ("exp2", ("log2", "x")), 1))
    g.run(l2.ALGEBRA)
    assert g.extract(root, l2.op_cost) == (0, "x")
    assert g.extract(root, l2.ast_size) == (1, "x")


def test_reassociation_is_not_bit_exact():
    env = {"a": 1e8, "b": -1e8, "c": 1.0}
    left = l2.evaluate(("+", ("+", "a", "b"), "c"), env, np.float32)
    right = l2.evaluate(("+", "a", ("+", "b", "c")), env, np.float32)
    assert left == 1.0 and right == 0.0      # equal in the e-graph, different in fp32

"""L3 · Search cost vs measured latency: a Luminal-style genome search — your implementation (run ./check.sh l3).

After saturation the e-graph holds many programs. Luminal's CUDA backend does not trust a cost
formula to choose between them; it *measures*. The pieces (crates/luminal_cuda_lite/src/search.rs,
finalists.rs, src/extraction.rs):

  Genome      one producer choice per e-class (which e-node implements this class)
  GA          a genetic search over genomes; each candidate is compiled + run, metric = measured ns
  rank_insert keep only the best `keep` genomes, sorted by measured ns
  finalists   walk the ranking; a candidate that fails validation (e.g. its arena exceeds the
              device memory budget) is rejected and the next rank is tried

Here the "e-graph" is a dict {class: [option, ...]} and `measure(genome)` is injectable.
`mlp_space()` is a small but honest model of y = relu(x @ w + b) on a GPU.
"""
import random
from itertools import product
from math import prod


def space_size(space):
    return prod(len(v) for v in space.values())


def all_genomes(space):
    keys = sorted(space)
    return [dict(zip(keys, combo)) for combo in product(*(space[k] for k in keys))]


def fingerprint(genome):
    return tuple(sorted(genome.items()))


def random_genome(space, rng):
    return {k: rng.choice(v) for k, v in sorted(space.items())}


def mutate(genome, space, rng, p=0.3):
    g = dict(genome)
    for k in sorted(space):
        if rng.random() < p:
            g[k] = rng.choice(space[k])
    return g


def crossover(a, b, rng):
    return {k: (a if rng.random() < 0.5 else b)[k] for k in sorted(a)}


def rank_insert(ranked, nanos, genome, keep):
    """Port of search.rs rank_insert: ranked is [(nanos, genome)] ascending, at most `keep` long."""
    raise NotImplementedError("rank_insert")


def genetic_search(space, measure, budget, population=6, keep=4, seed=0, mutation=0.3):
    """Measure at most `budget` distinct genomes. Returns ranked finalists, the number of distinct
    measurements, and the best-so-far curve (one entry per measurement)."""
    raise NotImplementedError("genetic_search")


def proxy_search(space, proxy, measure, k):
    """The cheap alternative: rank *every* genome by a cost proxy, measure only the top k."""
    raise NotImplementedError("proxy_search")


def spearman(xs, ys):
    """Rank correlation (no tie correction needed for the tests; ties get average rank)."""
    raise NotImplementedError("spearman")


def select_finalist(ranked, validate):
    """finalists.rs: walk ranks 1, 2, ...; return (rank, nanos, genome, rejected[]) of the first
    candidate that validates. validate(genome) -> None if OK, else a reason string."""
    raise NotImplementedError("select_finalist")


def break_even_runs(search_seconds_a, latency_a, search_seconds_b, latency_b):
    """How many executions until spending more search time (a) pays off over b. inf if never."""
    raise NotImplementedError("break_even_runs")


# ---------------------------------------------------------------- the workload
def mlp_space():
    """Decision points for y = relu(x @ w + b) after saturation (simplified)."""
    return {
        "matmul": ["primitive_loop", "tiled_kernel", "cublaslt"],
        "epilogue": ["separate", "fused"],
        "w_layout": ["row_major", "col_major"],
        "accumulate": ["fp32", "fp16"],
    }


MATMUL_NS = {"primitive_loop": 2400, "tiled_kernel": 130, "cublaslt": 95}
FUSE_SAVING = {"primitive_loop": 0, "tiled_kernel": 25, "cublaslt": 22}     # cuBLASLt: EPILOGUE_RELU_BIAS


def mlp_measure(noise=0.0, seed=0):
    """'Measured' ns: what a GPU would show, including effects a FLOP proxy cannot see."""
    rng = random.Random(seed)

    def m(g):
        ns = MATMUL_NS[g["matmul"]]
        ns += 30 if g["epilogue"] == "separate" else 30 - FUSE_SAVING[g["matmul"]]
        if g["matmul"] == "tiled_kernel" and g["w_layout"] == "row_major":
            ns += 20                      # B tile loads along K are strided for this kernel
        if g["accumulate"] == "fp16":
            ns -= 8
        return ns * (1 + rng.gauss(0, noise)) if noise else ns
    return m


def mlp_proxy(g):
    """A FLOP/byte proxy: knows the primitive loop is bad and fusion saves a pass, but cannot
    tell a hand-written tile from cuBLASLt, nor see layout effects."""
    base = {"primitive_loop": 2000, "tiled_kernel": 100, "cublaslt": 100}[g["matmul"]]
    return base + (30 if g["epilogue"] == "separate" else 10) + (-5 if g["accumulate"] == "fp16" else 0)


def mlp_validate(budget_bytes=8 << 20, M=64, N=64, K=256, max_err_ulps=None):
    """Reject plans whose intermediates exceed the arena budget, and fp16 accumulation (precision)."""
    def v(g):
        if g["matmul"] == "primitive_loop" and M * N * K * 4 > budget_bytes:
            return f"arena {M * N * K * 4} B > budget {budget_bytes} B"
        if g["accumulate"] == "fp16":
            return "fp16 accumulation fails the K=%d tolerance" % K
        return None
    return v

import pytest

from drill import load

l3 = load("l3_search")
S = l3.mlp_space()


def test_space():
    assert l3.space_size(S) == 3 * 2 * 2 * 2
    gs = l3.all_genomes(S)
    assert len(gs) == 24 and len({l3.fingerprint(g) for g in gs}) == 24


def test_rank_insert_keeps_best_k_sorted():
    r = []
    for ns, name in [(50, "a"), (30, "b"), (70, "c"), (10, "d"), (60, "e")]:
        l3.rank_insert(r, ns, name, keep=3)
    assert r == [(10, "d"), (30, "b"), (50, "a")]
    l3.rank_insert(r, 50, "tie", keep=3)        # cannot displace an equal worst
    assert r[-1] == (50, "a")


def test_true_optimum():
    m = l3.mlp_measure()
    best = min(l3.all_genomes(S), key=m)
    assert m(best) == 95 and best["matmul"] == "cublaslt" and best["epilogue"] == "fused"


def test_genetic_search_budget_and_curve():
    m = l3.mlp_measure()
    r = l3.genetic_search(S, m, budget=10, seed=0)
    assert r["measured"] == 10 and len(r["curve"]) == 10
    assert all(a >= b for a, b in zip(r["curve"], r["curve"][1:]))     # best-so-far never gets worse
    assert r["ranked"][0][0] == 95 and len(r["ranked"]) == 4
    small = l3.genetic_search(S, m, budget=3, seed=1)
    assert small["measured"] == 3


def test_proxy_is_correlated_but_not_exact():
    m = l3.mlp_measure()
    gs = l3.all_genomes(S)
    rho = l3.spearman([l3.mlp_proxy(g) for g in gs], [m(g) for g in gs])
    assert 0.7 < rho < 0.95
    assert l3.spearman([1, 2, 3], [10, 20, 30]) == pytest.approx(1.0)
    assert l3.spearman([1, 2, 3], [3, 2, 1]) == pytest.approx(-1.0)


def test_finalists_walk_past_invalid_candidates():
    m = l3.mlp_measure()
    r = l3.genetic_search(S, m, budget=24, seed=0)
    rank, ns, g, rejected = l3.select_finalist(r["ranked"], l3.mlp_validate())
    assert rank == 3 and ns == 103 and g["accumulate"] == "fp32" and len(rejected) == 2
    assert all("fp16" in why for _, why in rejected)


def test_proxy_top_k_can_all_be_invalid():
    ranked = l3.proxy_search(S, l3.mlp_proxy, l3.mlp_measure(), k=3)
    with pytest.raises(RuntimeError, match="no finalist"):
        l3.select_finalist(ranked, l3.mlp_validate())


def test_primitive_loop_breaks_the_arena_budget():
    v = l3.mlp_validate(budget_bytes=1 << 20)
    g = {"matmul": "primitive_loop", "epilogue": "separate", "w_layout": "row_major", "accumulate": "fp32"}
    assert "arena" in v(g)


def test_break_even():
    # 2 s of search to save 30 ns per call vs 0.1 s of search
    assert l3.break_even_runs(2.0, 87e-9, 0.1, 117e-9) == pytest.approx(1.9 / 30e-9)
    assert l3.break_even_runs(1.0, 120e-9, 0.1, 117e-9) == float("inf")

"""L4 · Backend trace: from the selected plan to GPU launches and one memory slab — your implementation (run ./check.sh l4).

Once search picked a genome, the backend turns the plan into:

  1. a DAG of kernels / library calls (the LLIR the README's LUMINAL_LOG_LLIR prints),
  2. buffer lifetimes: the DAG nodes at which each buffer is read or written,
  3. one device slab with an offset per buffer — src/buffer/arena.rs `plan_slab`:
       * transitive closure of the DAG ("u happens-before v"),
       * two buffers may share bytes only if *every* use of one precedes *every* use of the other,
       * first fit in request order, sizes rounded up to the alignment (256 B on CUDA),
  4. an execution timeline, split into phases (graph compile / search / prepare / execute —
     what LUMINAL_CUDA_PROFILE_RECAPTURE breaks TTFT/TPOT into).
"""
from dataclasses import dataclass


@dataclass
class SlabBuffer:
    id: str
    bytes: int
    uses: list          # DAG node names


def toposort(dag):
    """dag: {node: [successors]}. Kahn's algorithm with sorted tie-break (deterministic)."""
    raise NotImplementedError("toposort")


def reachability(dag):
    """after[u] = set of nodes strictly reachable from u."""
    raise NotImplementedError("reachability")


def aligned(n, a):
    return -(-max(n, 1) // a) * a          # empty buffers still reserve one aligned unit


def plan_slab(dag, buffers, capacity, alignment=256):
    """Port of arena.rs plan_slab (FirstFit). Returns {buffer id: offset}; raises on overflow."""
    raise NotImplementedError("plan_slab")


def slab_bytes(buffers, offsets, alignment=256):
    """High-water mark: what finalists.rs compares to the device budget."""
    raise NotImplementedError("slab_bytes")


def naive_bytes(buffers, alignment=256):
    """One allocation per buffer, no reuse."""
    raise NotImplementedError("naive_bytes")


def check_no_overlap(dag, buffers, offsets, alignment=256):
    """Invariant: buffers whose lifetimes are not ordered never share a byte."""
    raise NotImplementedError("check_no_overlap")


# ---------------------------------------------------------------- the MLP plan
def mlp_plan(M=64, K=256, N=64, fused=True, elem=4):
    """The selected plan for relu(x @ w + b) as (dag, buffers, launches).
    fused=True : one cuBLASLt call with the RELU_BIAS epilogue.
    fused=False: cuBLASLt matmul, then an add-bias kernel, then a relu kernel."""
    if fused:
        dag = {"upload": ["lt_matmul_relu_bias"], "lt_matmul_relu_bias": ["readback"], "readback": []}
        bufs = [SlabBuffer("x", M * K * elem, ["upload", "lt_matmul_relu_bias"]),
                SlabBuffer("w", K * N * elem, ["upload", "lt_matmul_relu_bias"]),
                SlabBuffer("b", N * elem, ["upload", "lt_matmul_relu_bias"]),
                SlabBuffer("workspace", 1 << 20, ["lt_matmul_relu_bias"]),
                SlabBuffer("y", M * N * elem, ["lt_matmul_relu_bias", "readback"])]
        launches = [("lt_matmul_relu_bias", "cublasLtMatmul", 2 * M * N * K)]
    else:
        dag = {"upload": ["lt_matmul"], "lt_matmul": ["add_bias"], "add_bias": ["relu"],
               "relu": ["readback"], "readback": []}
        bufs = [SlabBuffer("x", M * K * elem, ["upload", "lt_matmul"]),
                SlabBuffer("w", K * N * elem, ["upload", "lt_matmul"]),
                SlabBuffer("b", N * elem, ["upload", "add_bias"]),
                SlabBuffer("workspace", 1 << 20, ["lt_matmul"]),
                SlabBuffer("t0", M * N * elem, ["lt_matmul", "add_bias"]),
                SlabBuffer("t1", M * N * elem, ["add_bias", "relu"]),
                SlabBuffer("y", M * N * elem, ["relu", "readback"])]
        launches = [("lt_matmul", "cublasLtMatmul", 2 * M * N * K),
                    ("add_bias", "kernel", M * N), ("relu", "kernel", M * N)]
    return dag, bufs, launches


# ---------------------------------------------------------------- phases
PHASES = ("graph_compile", "search", "prepare", "execute")


def phase_breakdown(events):
    """events: [(phase, seconds)] possibly repeated. Returns {phase: total, 'share': {phase: frac}}."""
    raise NotImplementedError("phase_breakdown")


def achieved(flops, seconds, peak_flops):
    """Achieved FLOP/s and fraction of peak — compare theory with the profiled kernel time."""
    raise NotImplementedError("achieved")

import pytest

from drill import load

l4 = load("l4_plan_trace")
B = l4.SlabBuffer


def test_toposort_and_reachability():
    dag = {"a": ["b", "c"], "b": ["d"], "c": ["d"], "d": []}
    assert l4.toposort(dag) == ["a", "b", "c", "d"]
    after = l4.reachability(dag)
    assert after["a"] == {"b", "c", "d"} and after["b"] == {"d"} and after["d"] == set()
    with pytest.raises(ValueError):
        l4.toposort({"a": ["b"], "b": ["a"]})


def test_chain_reuses_memory():
    dag = {"a": ["b"], "b": ["c"], "c": []}
    bufs = [B("p", 1000, ["a"]), B("q", 100, ["b"]), B("r", 300, ["c"]), B("s", 10, ["a", "c"])]
    off = l4.plan_slab(dag, bufs, 10_000)
    assert off == {"p": 0, "q": 0, "r": 0, "s": 1024}       # p, q, r never live together
    assert l4.slab_bytes(bufs, off) == 1280
    assert l4.naive_bytes(bufs) == 1024 + 256 + 512 + 256


def test_unordered_branches_cannot_alias():
    dag = {"a": ["b", "c"], "b": ["d"], "c": ["d"], "d": []}
    bufs = [B("p", 256, ["b"]), B("q", 256, ["c"]), B("r", 256, ["d"])]
    off = l4.plan_slab(dag, bufs, 10_000)
    assert off["p"] != off["q"]                              # b and c may run concurrently
    assert off["r"] == 0                                     # d is after both
    assert l4.check_no_overlap(dag, bufs, off)


def test_alignment_and_capacity():
    dag = {"a": []}
    assert l4.aligned(1, 256) == 256 and l4.aligned(0, 256) == 256 and l4.aligned(257, 256) == 512
    with pytest.raises(MemoryError, match="capacity"):
        l4.plan_slab(dag, [B("x", 300, ["a"]), B("y", 300, ["a"])], capacity=768)
    with pytest.raises(ValueError):
        l4.plan_slab(dag, [B("x", 1, ["a"])], 100, alignment=0)


def test_mlp_plans():
    dag, bufs, launches = l4.mlp_plan(fused=True)
    off = l4.plan_slab(dag, bufs, 64 << 20)
    assert len(launches) == 1 and launches[0][1] == "cublasLtMatmul"
    assert l4.check_no_overlap(dag, bufs, off)
    dag2, bufs2, launches2 = l4.mlp_plan(fused=False)
    off2 = l4.plan_slab(dag2, bufs2, 64 << 20)
    assert [l[0] for l in launches2] == ["lt_matmul", "add_bias", "relu"]
    assert off2["t1"] == 0 and off2["y"] == 16384            # x is dead by then: its bytes are reused
    assert l4.slab_bytes(bufs2, off2) < l4.naive_bytes(bufs2)
    assert l4.check_no_overlap(dag2, bufs2, off2)


def test_phases_and_achieved():
    ph = l4.phase_breakdown([("graph_compile", 0.4), ("search", 3.0), ("prepare", 0.1),
                             ("execute", 0.0005), ("execute", 0.0005)])
    assert ph["execute"] == pytest.approx(0.001)
    assert ph["share"]["search"] == pytest.approx(3.0 / 3.501)
    with pytest.raises(ValueError):
        l4.phase_breakdown([("lunch", 1.0)])
    f, frac = l4.achieved(2 * 64 * 64 * 256, 10e-6, 40e12)
    assert f == pytest.approx(2 * 64 * 64 * 256 / 10e-6) and frac < 0.01   # tiny GEMM: launch-bound

"""L5 · Differential test: ReferenceRuntime vs CUDA, written as a regression test — your implementation (run ./check.sh l5).

Luminal ships two runtimes for the same compiled graph: `luminal_reference` (a CPU interpreter
that defines the semantics, with explicit narrow-float and tf32 emulation: narrow_float.rs,
tf32.rs) and `luminal_cuda_lite` (the real device). A backend contribution is judged by
"does CUDA agree with Reference within a tolerance I can justify?".

This drill builds that harness in Python around one op (softmax over the last axis, a classic
place for backend bugs: max-subtraction, masked tails, fp16 overflow):

  reference_softmax   float64, the semantics
  device_softmax      a 'device' implementation: tiled in blocks, fp32 accumulation, output in
                      the storage dtype — with switchable, realistic bugs
  tolerance_for       rtol/atol from the dtype's eps and the reduction length (K3), plus a
                      documented rationale string that goes into the test
  run_differential    random shapes (edge-heavy), dtypes, seeds -> list of failures
  emit_rust_test      the regression test in the style of crates/luminal_cuda_lite/tests
"""
import math
import random

import numpy as np

from solutions.k3_numerics import EPS, check_close, edge_sizes, to_bf16

STORAGE = {"float32": np.float32, "float16": np.float16, "bfloat16": np.float32}


def store(x, dtype):
    x = np.asarray(x, np.float32).astype(STORAGE[dtype])
    return to_bf16(x) if dtype == "bfloat16" else x


def reference_softmax(x):
    x = np.asarray(x, np.float64)
    m = x.max(axis=-1, keepdims=True)
    e = np.exp(x - m)
    return e / e.sum(axis=-1, keepdims=True)


def device_softmax(x, dtype, block=64, bug=None):
    """Two passes over the row in blocks of `block`: online max+sum, then normalize.
    bug: None | 'no_max_sub' (overflows for large inputs) | 'tail' (last partial block ignored in
    the sum) | 'fp16_acc' (accumulate in the storage dtype)."""
    x = store(x, dtype).astype(np.float32)
    rows, n = x.shape
    acc_t = np.float16 if bug == "fp16_acc" else np.float32
    out = np.empty_like(x)
    for r in range(rows):
        m = np.float32(-np.inf) if bug != "no_max_sub" else np.float32(0)
        s = acc_t(0)
        stop = n if bug != "tail" else (n // block) * block or n
        for k0 in range(0, stop, block):
            blk = x[r, k0:min(k0 + block, stop)]
            if bug != "no_max_sub":
                m_new = max(m, blk.max())
                s = acc_t(s * acc_t(np.exp(m - m_new)) + acc_t(np.exp(blk - m_new).sum()))
                m = m_new
            else:
                with np.errstate(over="ignore"):
                    s = acc_t(s + acc_t(np.exp(blk).sum()))
        with np.errstate(over="ignore", invalid="ignore"):
            out[r] = np.exp(x[r] - m) / np.float32(s)
    return store(out, dtype)


def tolerance_for(dtype, n):
    """(rtol, atol, rationale). Softmax outputs are in [0, 1]; the error comes from (a) rounding
    the *output* to the storage dtype (half an ulp: eps/2 relative) and (b) the length-n sum in
    fp32 (~sqrt(n) * eps32). atol guards the tiny probabilities where relative error is noise."""
    raise NotImplementedError("tolerance_for")


def make_input(rows, n, dtype, seed, scale=3.0):
    rng = np.random.default_rng(seed)
    x = rng.standard_normal((rows, n)) * scale
    return store(x, dtype)


def run_differential(impl, cases=60, seed=0, block=64, dtypes=("float32", "float16", "bfloat16"), scale=3.0):
    """Returns (failures, tried). failures: [{'rows','n','dtype','seed','report'}]."""
    raise NotImplementedError("run_differential")


def smallest_failure(fails):
    raise NotImplementedError("smallest_failure")


RUST_DTYPE = {"float32": "F32", "float16": "F16", "bfloat16": "Bf16"}


def emit_rust_test(f, name="softmax_matches_reference"):
    """A Rust regression test sketch: same graph compiled for both runtimes, compared with the
    justified tolerance. (Adapt the API calls to the current crate before submitting.)"""
    raise NotImplementedError("emit_rust_test")

import numpy as np

from drill import load

l5 = load("l5_differential")


def impl(bug=None):
    return lambda x, d: l5.device_softmax(x, d, bug=bug)


def test_reference_is_stable():
    y = l5.reference_softmax(np.array([[1000.0, 1000.0], [0.0, -np.inf]]))
    np.testing.assert_allclose(y, [[0.5, 0.5], [1.0, 0.0]])


def test_correct_device_passes_everywhere():
    fails, tried = l5.run_differential(impl(None), cases=60, seed=0)
    assert tried == 60 and fails == []
    fails, _ = l5.run_differential(impl(None), cases=30, seed=1, scale=40.0)
    assert fails == []


def test_tail_bug_found_and_minimized():
    fails, _ = l5.run_differential(impl("tail"), cases=60, seed=0)
    assert len(fails) == 26
    f = l5.smallest_failure(fails)
    assert f["rows"] == 1 and f["n"] % 64 != 0 and f["n"] > 64


def test_fp16_accumulation_bug_found():
    fails, _ = l5.run_differential(impl("fp16_acc"), cases=60, seed=0)
    assert len(fails) == 11 and all(f["report"]["max_abs"] < 1e-2 for f in fails)   # small but real


def test_overflow_bug_needs_large_inputs():
    # N(0, 3) inputs never overflow exp(): the missing max-subtraction is invisible ...
    fails, _ = l5.run_differential(impl("no_max_sub"), cases=60, seed=0)
    assert fails == []
    # ... until the test includes large logits
    fails, _ = l5.run_differential(impl("no_max_sub"), cases=60, seed=0, scale=40.0)
    assert len(fails) == 44


def test_tolerance_rationale():
    rtol, atol, why = l5.tolerance_for("float16", 65)
    assert rtol > 2 * 2 ** -10 and atol < rtol
    assert "float16" in why and "n=65" in why
    r32, _, _ = l5.tolerance_for("float32", 65)
    assert r32 < rtol < l5.tolerance_for("bfloat16", 65)[0]


def test_emit_rust_test():
    fails, _ = l5.run_differential(impl("tail"), cases=60, seed=0)
    src = l5.emit_rust_test(l5.smallest_failure(fails))
    assert src.startswith("#[test]") and "softmax(1)" in src and "DType::F32" in src
    assert "tolerance:" in src and src.count("{") == src.count("}")