docs: regenerate the benchmark figures and add their scripts

.gitignore

@@ -11,3 +11,4 @@ build/
 .vscode/
 .venv/
 benchmarks/results/
+/.worktrees/

README.md

@@ -162,39 +162,37 @@ One exception is deliberate. The `training` flag stays a plain runtime boolean,
 
 ## Benchmarks
 
-All figures below come from the same CNN, the same CIFAR-10 batches, the same seed, and the same hardware (AMD RX 9070 XT, Intel Core Ultra 265K). Only the dispatch path differs.
+All figures below come from the same CNN, the same seed, the same fp16 precision policy, and the same hardware (AMD RX 9070 XT, Intel Core Ultra 265K). The dispatch path differs, and with it Conv2d's precision: only the fused path can run Conv2d in fp16, so the other two compute it in fp32. The scripts and saved results are in [`benchmarks/`](benchmarks/README.md).
 
 ### Convergence
-Four paths are tested; NumPy on CPU, legacy pre-refactor GPU implementation, vectorized CuPy without custom kernels, and the fused kernel path.  
+Three paths are tested; NumPy on CPU, vectorized CuPy without custom kernels, and the fused kernel path.  
 
 <picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_convergence-dark.png">
-  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_convergence.png">
-  <img alt="Backend convergence" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_convergence.png">
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_convergence-dark.png">
+  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_convergence.png">
+  <img alt="Backend convergence" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_convergence.png">
 </picture>
 
 ### Throughput
 
-Per training step: **3.17 ms** on the fused kernel path vs **52.6 ms** on vectorized CuPy on the same GPU, with **326 ms** for NumPy on CPU. With the GPU paths, the kernel work shows a **17x speedup against identical hardware running array operations**.
+Per training step: **4.85 ms** on the fused kernel path vs **202 ms** on vectorized CuPy on the same GPU, with **1,865 ms** for NumPy on CPU. With the GPU paths, the kernel work shows a **42x speedup against identical hardware running array operations**.
 
 <picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_throughput-dark.png">
-  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_throughput.png">
-  <img alt="Backend throughput" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_throughput.png">
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_throughput-dark.png">
+  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_throughput.png">
+  <img alt="Backend throughput" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_throughput.png">
 </picture>
 
 ### Transient VRAM 
 
-Transient VRAM footprint over two training steps. Unfused CuPy dispatches allocate independent temporary buffers in VRAM for intermediate steps, resulting in transient spikes up to **180 MiB**. Custom fused kernels eliminate the need to create temporary buffers by executing multiple actions in the same kernel within registers and shared memory, reducing peak transient VRAM by **5.6×** (32.0 MiB peak).  
+Transient VRAM footprint over two training steps. Unfused CuPy dispatches allocate independent temporary buffers in VRAM for intermediate steps, resulting in transient spikes up to **420 MiB**. Custom fused kernels eliminate the need to create temporary buffers by executing multiple actions in the same kernel within registers and shared memory, reducing peak transient VRAM by **5.3×** (80.0 MiB peak).  
 
 <picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_vram-dark.png">
-  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_vram.png">
-  <img alt="Transient VRAM usage" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmark_vram.png">
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_vram-dark.png">
+  <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_vram.png">
+  <img alt="Transient VRAM usage" src="https://raw.githubusercontent.com/AlexanderSoftCode/Aether-ML/main/notebooks/assets/benchmarks/benchmark_vram.png">
 </picture>
 
-**On the legacy series.** "Legacy CuPy" contains a `float64` type promotion bug that was found during the rewrite. GPUs often have very few or no FP64 ALUs, which explains the slowdown. The pathway is still included as an honest before-and-after, and not as a competitive baseline.
-
 ---
 
 ## Building from Source

benchmarks/README.md

@@ -0,0 +1,61 @@
+# Benchmarks
+
+The figures in the main README and the release notes come from the scripts here, on one
+CNN: `build_model` in `cifar10_breakdown.py`, the same model as
+`examples/cifar10/v1_cifar10.py`.
+
+| Script | What it measures | Figures |
+|---|---|---|
+| `cifar10_breakdown.py` | ms/step, split by training-step phase, on the normal GPU path | `benchmark_release` |
+| `cifar10_paths.py` | ms/step and loss over one epoch, on NumPy, fallback CuPy and fused kernels | `benchmark_throughput`, `benchmark_convergence` |
+| `cifar10_vram.py` | live VRAM over two steps, on fallback CuPy and fused kernels | `benchmark_vram` |
+| `compare.py` | prints what changed between two `cifar10_breakdown.py` results | none |
+| `plot.py` | draws the figures from saved results | all of the above |
+
+The scripts need a GPU and CuPy. They also need CIFAR-10 in `data/cifar-10`: create it with
+`examples/cifar10/cifar-setup.ipynb`. `plot.py` needs matplotlib instead, which isn't a
+dependency of the package.
+
+## Before and after a change
+
+    python3 benchmarks/cifar10_breakdown.py benchmarks/results/before.json
+    python3 benchmarks/cifar10_breakdown.py benchmarks/results/after.json
+    python3 benchmarks/compare.py benchmarks/results/before.json benchmarks/results/after.json
+
+`benchmarks/results/` is ignored by git, so it's the place for working numbers.
+
+## Regenerating the figures for a release
+
+Published numbers live in `benchmarks/history/<version>/`, so a figure can be redrawn
+without rerunning anything. Measure from clean worktrees: the scripts stay in your
+checkout, `PYTHONPATH` picks the `aether` being measured, and each result records that
+tree's commit. `--relative-paths` lets the worktree resolve inside the devcontainer.
+
+    git worktree add --relative-paths --detach .worktrees/old v1.1.0
+    git worktree add --relative-paths --detach .worktrees/new <release commit>
+
+Run from the repo root on each GPU, naming the file after that GPU:
+
+    PYTHONPATH=.worktrees/old python3 benchmarks/cifar10_breakdown.py benchmarks/history/v1.1.0/breakdown_rtx3050.json
+    PYTHONPATH=.worktrees/new python3 benchmarks/cifar10_breakdown.py benchmarks/history/v1.2.0/breakdown_rtx3050.json
+
+On a GPU other than the one the README names, also add its fused path to the convergence
+figure:
+
+    PYTHONPATH=.worktrees/new python3 benchmarks/cifar10_paths.py --fused-only benchmarks/history/v1.2.0/paths_fused_rtx3050.json
+
+Then once, on the machine the README names, for the path figures:
+
+    PYTHONPATH=.worktrees/new python3 benchmarks/cifar10_vram.py benchmarks/history/v1.2.0/vram.json
+    PYTHONPATH=.worktrees/new python3 benchmarks/cifar10_paths.py benchmarks/history/v1.2.0/paths.json
+
+Draw them, then remove the worktrees:
+
+    python3 benchmarks/plot.py v1.1.0 v1.2.0
+    git worktree remove .worktrees/old
+    git worktree remove .worktrees/new
+
+`cifar10_paths.py` runs NumPy once and each GPU path three times, which takes about 15
+minutes. The cross-version figure uses only `cifar10_breakdown.py`, so older versions need nothing
+else from this folder. Each GPU gets its own panel and scale on it: comparing versions is
+only a fair A/B on one machine.

benchmarks/cifar10_breakdown.py

@@ -36,8 +36,10 @@
 from pathlib import Path
 
 import cupy as cp
+import numpy as np
 
 import aether as ae
+import aether.config as config
 
 TRAIN_DIR = Path("data") / "cifar-10" / "cifar-10_train.npz"
 
@@ -48,10 +50,22 @@
 BREAKDOWN_STEPS = 100
 
 
-def build_model(X_train):
+def load_train():
+    """CIFAR-10 training images and labels, as host arrays."""
+    if not TRAIN_DIR.exists():
+        sys.exit(
+            f"[aether] {TRAIN_DIR} not found. Run from the repo root after creating it "
+            "with examples/cifar10/cifar-setup.ipynb."
+        )
+    with np.load(TRAIN_DIR, allow_pickle=False) as data:
+        return data["X_train"], data["y_train"]
+
+
+def build_model(X_train, device="cupy"):
     """Same model, pipeline and precision as examples/cifar10/v1_cifar10.py."""
+    config.set_backend(device)
     feature_pipeline = ae.Compose([
-        ae.ToTensor(dtype='float32', target_device="cupy"),
+        ae.ToTensor(dtype='float32', target_device=device),
         ae.Rescale(factor=1.0 / 255.0),
         ae.StandardScaler()
     ])
@@ -93,7 +107,7 @@ def build_model(X_train):
         preprocessor=feature_pipeline
     )
 
-    model.to('cupy')
+    model.to(device)
     model.set_precision(compute_dtype=COMPUTE_DTYPE)
     model.finalize(input_shape=X_train.shape[1:])
     feature_pipeline.fit(X_train)
@@ -218,17 +232,36 @@ def spread(samples):
 
 
 def git_state():
-    """(commit, dirty) of the checkout, or (None, None) outside git."""
+    """(commit, dirty) of the checkout aether was imported from, or (None, None)
+    outside git, so a run against a worktree of an older tag records that tag."""
+    source = Path(ae.__file__).parent
     try:
         commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True,
-                                text=True, check=True).stdout.strip()
+                                text=True, check=True, cwd=source).stdout.strip()
         status = subprocess.run(["git", "status", "--porcelain"], capture_output=True,
-                                text=True, check=True).stdout
+                                text=True, check=True, cwd=source).stdout
     except (OSError, subprocess.CalledProcessError):
         return None, None
     return commit, bool(status.strip())
 
 
+def run_meta():
+    """What a result file needs to say about where it was measured."""
+    commit, dirty = git_state()
+    return {
+        "commit": commit,
+        "dirty": dirty,
+        "date": datetime.datetime.now().isoformat(timespec="seconds"),
+        "gpu": cp.cuda.runtime.getDeviceProperties(0)["name"].decode(),
+        "cupy": cp.__version__,
+        "python": platform.python_version(),
+        "batch_size": BATCH_SIZE,
+        "compute_dtype": COMPUTE_DTYPE,
+        "warmup_steps": WARMUP_STEPS,
+        "baseline_steps": BASELINE_STEPS,
+    }
+
+
 def bound_paths(model):
     """Forward function each layer is bound to, keyed like '03 Conv2d'."""
     paths = {}
@@ -249,14 +282,7 @@ def main():
     if args.repeats < 1:
         parser.error("--repeats must be at least 1")
 
-    if not TRAIN_DIR.exists():
-        sys.exit(
-            f"[aether] {TRAIN_DIR} not found. Run from the repo root after creating it "
-            "with examples/cifar10/cifar-setup.ipynb."
-        )
-    with cp.load(TRAIN_DIR, allow_pickle=False) as data:
-        X_train = data["X_train"]
-        y_train = data["y_train"]
+    X_train, y_train = (cp.asarray(array) for array in load_train())
 
     model, feature_pipeline = build_model(X_train)
     transforms = feature_pipeline.transforms
@@ -271,19 +297,9 @@ def main():
             timed_train(model, X_train, y_train, BASELINE_STEPS) * 1e3 / BASELINE_STEPS)
         breakdown_runs.append(breakdown(model, transforms, X_train, y_train))
 
-    commit, dirty = git_state()
     result = {
         "meta": {
-            "commit": commit,
-            "dirty": dirty,
-            "date": datetime.datetime.now().isoformat(timespec="seconds"),
-            "gpu": cp.cuda.runtime.getDeviceProperties(0)["name"].decode(),
-            "cupy": cp.__version__,
-            "python": platform.python_version(),
-            "batch_size": BATCH_SIZE,
-            "compute_dtype": COMPUTE_DTYPE,
-            "warmup_steps": WARMUP_STEPS,
-            "baseline_steps": BASELINE_STEPS,
+            **run_meta(),
             "breakdown_steps": BREAKDOWN_STEPS,
             "repeats": args.repeats,
             "warmup_s": warmup_s,

benchmarks/cifar10_paths.py

@@ -0,0 +1,155 @@
+"""
+Train the same CIFAR-10 CNN on three dispatch paths, for the README's throughput and
+convergence figures.
+
+Run from the repo root (see benchmarks/README.md for the full release workflow):
+
+    python3 benchmarks/cifar10_paths.py benchmarks/history/v1.2.0/paths.json
+
+--fused-only skips the slower paths, for adding another GPU to the convergence figure.
+
+The paths, all built by cifar10_breakdown.build_model:
+
+    numpy-cpu      the framework on the CPU fallback path
+    cupy-fallback  tensors on the GPU, but every layer and the optimizer on that same
+                   fallback path, so vectorized CuPy array ops and no custom kernels
+    cupy-fused     the normal model.to("cupy") path
+
+The precision policy is the same everywhere, but only the fused path's Conv2d can honor
+fp16; the fallback paths warn and compute Conv2d in fp32.
+
+Each run warms up, records loss and elapsed time for one epoch, then times BASELINE_STEPS
+with no per-step sync. The epoch syncs after every step so each loss gets a timestamp,
+which costs time on a fast step, so throughput comes from the unsynced block. The first
+full-length train() runs slow (see cifar10_breakdown.main), and the epoch absorbs that.
+"""
+import argparse
+import gc
+import json
+import time
+from pathlib import Path
+
+import cupy as cp
+from cifar10_breakdown import (
+    BASELINE_STEPS,
+    BATCH_SIZE,
+    WARMUP_STEPS,
+    bound_paths,
+    build_model,
+    load_train,
+    run_meta,
+    spread,
+    timed_train,
+)
+
+LABELS = {
+    "numpy-cpu": "NumPy CPU",
+    "cupy-fallback": "CuPy, no custom kernels",
+    "cupy-fused": "CuPy, fused kernels",
+}
+NUMPY_REPEATS = 1  # about 11 minutes a run, and the CPU is the least noisy path
+
+
+def use_fallback_kernels(model):
+    """Point a finalized GPU model's dispatch at the fallback path, leaving tensors on
+    the GPU. The fallbacks take their array module from the data, so they run as CuPy."""
+    for component in (*model.layers, model.loss, model.accuracy):
+        if hasattr(component, "_compile_for_device"):
+            component._compile_for_device("numpy")
+    # Not Adam's _compile_for_device, which also moves its buffers to the host.
+    model.optimizer._step_impl = model.optimizer._step_fallback
+    # BatchNorm just lost its fused activation, so the skipped ReLUs have to come back.
+    model._bind_fusion()
+
+
+def prepare(key, X):
+    model, _ = build_model(X, device="numpy" if key == "numpy-cpu" else "cupy")
+    if key == "cupy-fallback":
+        use_fallback_kernels(model)
+    return model
+
+
+def trace_epoch(model, X, y):
+    """Per-step data loss and cumulative seconds over one epoch of the real train()."""
+    calculate, step_optimizer = model.loss.calculate, model._step_optimizer
+    pending_loss = None
+    losses, elapsed_s = [], []
+    total = 0.0
+
+    def remember_loss(*args, **kwargs):
+        nonlocal pending_loss
+        pending_loss = calculate(*args, **kwargs)
+        return pending_loss
+
+    def end_step():
+        nonlocal total, step_start
+        step_optimizer()
+        cp.cuda.Device().synchronize()
+        total += time.perf_counter() - step_start
+        elapsed_s.append(total)
+        losses.append(float(pending_loss))
+        step_start = time.perf_counter()  # the host readback above isn't step time
+
+    model.loss.calculate = remember_loss
+    model._step_optimizer = end_step
+    try:
+        step_start = time.perf_counter()
+        model.train(X, y, epochs=1, batch_size=BATCH_SIZE, shuffle=True, verbose=0,
+                    fit_preprocessor=False)
+    finally:
+        del model.loss.calculate  # it was a plain class method
+        model._step_optimizer = step_optimizer
+    return losses, elapsed_s
+
+
+def run_path(key, X, y, repeats):
+    steps_per_epoch = len(X) // BATCH_SIZE
+    epoch_X, epoch_y = X[:steps_per_epoch * BATCH_SIZE], y[:steps_per_epoch * BATCH_SIZE]
+    ms_per_step = []
+    for _ in range(repeats):
+        model = prepare(key, X)
+        timed_train(model, X, y, WARMUP_STEPS)
+        losses, elapsed_s = trace_epoch(model, epoch_X, epoch_y)
+        ms_per_step.append(timed_train(model, X, y, BASELINE_STEPS) * 1e3 / BASELINE_STEPS)
+        paths = bound_paths(model)
+        del model
+        gc.collect()
+        cp.get_default_memory_pool().free_all_blocks()
+    return {
+        "key": key,
+        "label": LABELS[key],
+        "repeats": repeats,
+        "ms_per_step": spread(ms_per_step),
+        "loss": losses,  # from the last repeat
+        "elapsed_s": elapsed_s,
+        "paths": paths,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[1])
+    parser.add_argument("output", type=Path, help="results file, for benchmarks/plot.py")
+    parser.add_argument("--repeats", type=int, default=3,
+                        help=f"runs per GPU path (default: 3); NumPy runs {NUMPY_REPEATS}")
+    parser.add_argument("--fused-only", action="store_true",
+                        help="run only cupy-fused, for a second GPU on the convergence figure")
+    args = parser.parse_args()
+    if args.repeats < 1:
+        parser.error("--repeats must be at least 1")
+
+    X, y = load_train()
+    X_gpu, y_gpu = cp.asarray(X), cp.asarray(y)
+    results = [run_path("cupy-fused", X_gpu, y_gpu, args.repeats)]
+    if not args.fused_only:
+        results = [
+            run_path("numpy-cpu", X, y, NUMPY_REPEATS),
+            run_path("cupy-fallback", X_gpu, y_gpu, args.repeats),
+            *results,
+        ]
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps({"meta": run_meta(), "results": results}, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()