feat(benchmarks): add a compiled PyTorch reference bar to the release figure

benchmarks/README.md

@@ -9,13 +9,21 @@ CNN: `build_model` in `cifar10_breakdown.py`, the same model as
 | `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` |
+| `cifar10_torch.py` | ms/step of the same CNN in PyTorch, compiled with fp16 autocast | `benchmark_release` |
 | `compare.py` | prints what changed between two `cifar10_breakdown.py` results | none |
+| `common.py` | the constants and helpers the scripts share; imports neither CuPy nor PyTorch | 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.
 
+`cifar10_torch.py` needs PyTorch, also not a dependency: install the build for your GPU from
+pytorch.org (on a ROCm GPU, the wheel index for your ROCm version). The result records the
+PyTorch version. PyTorch and CuPy can't be loaded into one process (`import cupy, torch`
+aborts with an LLVM error), so the torch script runs on its own and the figure joins the two
+result files.
+
 ## Before and after a change
 
     python3 benchmarks/cifar10_breakdown.py benchmarks/results/before.json
@@ -39,6 +47,13 @@ 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
 
+Then PyTorch's bar for the same GPU, from `torch_<gpu>.json` in the new version's folder. It
+is the same CNN in PyTorch with the whole step under `torch.compile`, fp16 autocast and no
+loss scaling, so it is a reference point and not a tuned ceiling. Without the file, the figure
+draws that GPU's panel without the bar:
+
+    python3 benchmarks/cifar10_torch.py benchmarks/history/v1.2.0/torch_rtx3050.json
+
 On a GPU other than the one the README names, also add its fused path to the convergence
 figure:
 
@@ -56,6 +71,8 @@ Draw them, then remove the worktrees:
     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
+minutes. The cross-version figure uses only `cifar10_breakdown.py` (and `cifar10_torch.py` for the
+PyTorch bar), so older versions need nothing else from this folder. The PyTorch bar is
+wall-clock ms/step, while the other bars are GPU time per stage, which runs a few percent above
+their wall-clock total. 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

@@ -26,41 +26,27 @@
 it does outside the wrapped methods grows "unattributed".
 """
 import argparse
-import datetime
 import json
-import platform
-import statistics
-import subprocess
-import sys
 import time
 from pathlib import Path
 
 import cupy as cp
-import numpy as np
+from common import (
+    BASELINE_STEPS,
+    BATCH_SIZE,
+    COMPUTE_DTYPE,
+    WARMUP_STEPS,
+    base_meta,
+    load_train,
+    spread,
+)
 
 import aether as ae
 import aether.config as config
 
-TRAIN_DIR = Path("data") / "cifar-10" / "cifar-10_train.npz"
-
-BATCH_SIZE = 256
-COMPUTE_DTYPE = "float16"
-WARMUP_STEPS = 20
-BASELINE_STEPS = 128
 BREAKDOWN_STEPS = 100
 
 
-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)
@@ -227,38 +213,12 @@ def timed(*args, **kwargs):
     return result
 
 
-def spread(samples):
-    return {"median": statistics.median(samples), "min": min(samples), "max": max(samples)}
-
-
-def git_state():
-    """(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, cwd=source).stdout.strip()
-        status = subprocess.run(["git", "status", "--porcelain"], capture_output=True,
-                                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"),
+        **base_meta(Path(ae.__file__).parent),
         "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,
     }
 
 

benchmarks/cifar10_torch.py

@@ -0,0 +1,143 @@
+"""
+Train the v1 CIFAR-10 CNN in PyTorch, as the reference bar on the release figure.
+
+Run from the repo root, in its own process: PyTorch and CuPy can't share one (see
+benchmarks/README.md). PyTorch isn't a dependency of the package; install the build for your
+GPU from pytorch.org.
+
+    python3 benchmarks/cifar10_torch.py benchmarks/history/v1.2.0/torch_rx9070xt.json
+
+The model is build_model in cifar10_breakdown.py, written the way a PyTorch user would, with
+the whole step under torch.compile: default NCHW layout, fp16 autocast and no loss scaling.
+Autocast also keeps the activations after each Conv2d in fp16, where Aether keeps them fp32.
+Every step does what Aether's does: standardize the batch on the GPU, forward, label-smoothed
+loss plus the L2 term on the third Conv2d, accuracy, backward and an AdamW step. Aether's
+learning-rate decay is left out, because it changes the loss and not the time per step.
+
+The timing protocol is cifar10_breakdown.py's baseline: warm up, discard one full-length run,
+then time BASELINE_STEPS per repeat with no per-step sync. Compilation happens in the warm-up,
+so it counts toward warmup_s and not baseline_ms.
+"""
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+from common import (
+    BASELINE_STEPS,
+    BATCH_SIZE,
+    WARMUP_STEPS,
+    base_meta,
+    load_train,
+    spread,
+)
+
+try:
+    import torch
+    import torch.nn.functional as F
+    from torch import nn
+except ImportError:
+    sys.exit("[aether] PyTorch is not installed; install the build for your GPU from pytorch.org.")
+
+SEED = 42  # build_model's manual_seed
+LABEL_SMOOTHING = 0.05
+L2 = 1e-5  # on the third Conv2d only
+PARAMETERS = 141_354  # 3 Conv2d, 5 BatchNorm (gamma, beta) and the Dense, counted by hand
+
+
+def build_model():
+    def conv_block(in_channels, out_channels):
+        return [nn.Conv2d(in_channels, out_channels, 3, padding="same"),
+                nn.BatchNorm2d(out_channels, eps=1e-5, momentum=0.1),  # Aether's momentum is 0.9
+                nn.ReLU()]
+
+    model = nn.Sequential(
+        *conv_block(3, 32), *conv_block(32, 32), nn.MaxPool2d(2, 2), nn.Dropout2d(0.1),
+        *conv_block(32, 64), *conv_block(64, 64), nn.MaxPool2d(2, 2), nn.Dropout2d(0.15),
+        *conv_block(64, 128), nn.MaxPool2d(2, 2), nn.Dropout2d(0.2),
+        nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(128, 10),
+    ).cuda()
+    parameters = sum(p.numel() for p in model.parameters())
+    if parameters != PARAMETERS:
+        sys.exit(f"[aether] The PyTorch model has {parameters} parameters, not {PARAMETERS}; "
+                 "it no longer matches build_model in cifar10_breakdown.py.")
+    return model
+
+
+def make_step(X):
+    """A compiled training step on the GPU."""
+    torch.manual_seed(SEED)
+    model = build_model()
+    l2_weight = model[8].weight  # the third Conv2d
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
+    scaled = X.float() / 255.0
+    mean, std = scaled.mean(), scaled.std(correction=0)
+    del scaled
+    correct = torch.zeros((), device=X.device, dtype=torch.int64)
+
+    def step(x, y):
+        x = ((x.float() * (1.0 / 255.0) - mean) / std).permute(0, 3, 1, 2).contiguous()
+        with torch.autocast("cuda", dtype=torch.float16):
+            logits = model(x)
+        logits = logits.float()
+        loss = (F.cross_entropy(logits, y, label_smoothing=LABEL_SMOOTHING)
+                + L2 * l2_weight.square().sum())
+        correct.add_((logits.argmax(1) == y).sum())
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        return loss
+
+    return torch.compile(step)
+
+
+def timed_train(step, X, y, num_steps):
+    """Wall time of num_steps shuffled full batches."""
+    n = num_steps * BATCH_SIZE
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    order = torch.randperm(n, device=X.device)
+    for i in range(num_steps):
+        batch = order[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
+        step(X[batch], y[batch])
+    torch.cuda.synchronize()
+    return time.perf_counter() - start
+
+
+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="baseline runs after warm-up (default: 3)")
+    args = parser.parse_args()
+    if args.repeats < 1:
+        parser.error("--repeats must be at least 1")
+    if not torch.cuda.is_available():
+        sys.exit("[aether] PyTorch can't see a GPU.")
+
+    X, y = (torch.as_tensor(array).cuda() for array in load_train())
+    step = make_step(X)
+
+    warmup_s = timed_train(step, X, y, WARMUP_STEPS)
+    timed_train(step, X, y, BASELINE_STEPS)  # the first full-length run is slow
+    baseline_runs = [timed_train(step, X, y, BASELINE_STEPS) * 1e3 / BASELINE_STEPS
+                     for _ in range(args.repeats)]
+
+    result = {
+        "meta": {
+            **base_meta(Path(__file__).resolve().parent),
+            "gpu": torch.cuda.get_device_name(0),
+            "torch": torch.__version__,
+            "mode": "compiled",
+            "repeats": args.repeats,
+            "warmup_s": warmup_s,
+        },
+        "baseline_ms": spread(baseline_runs),
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(result, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/common.py

@@ -0,0 +1,63 @@
+"""
+What the benchmark scripts share. It imports neither CuPy nor PyTorch, because the two
+can't be loaded into one process (see benchmarks/README.md).
+"""
+import datetime
+import platform
+import statistics
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+
+TRAIN_DIR = Path("data") / "cifar-10" / "cifar-10_train.npz"
+
+BATCH_SIZE = 256
+COMPUTE_DTYPE = "float16"
+WARMUP_STEPS = 20
+BASELINE_STEPS = 128
+
+
+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 spread(samples):
+    return {"median": statistics.median(samples), "min": min(samples), "max": max(samples)}
+
+
+def git_state(source):
+    """(commit, dirty) of the checkout containing `source`, or (None, None) outside git,
+    so a run against a worktree of an older tag records that tag."""
+    try:
+        commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True,
+                                text=True, check=True, cwd=source).stdout.strip()
+        status = subprocess.run(["git", "status", "--porcelain"], capture_output=True,
+                                text=True, check=True, cwd=source).stdout
+    except (OSError, subprocess.CalledProcessError):
+        return None, None
+    return commit, bool(status.strip())
+
+
+def base_meta(source):
+    """What a result file says about the run, whatever library ran it. `source` is the
+    directory of the code that was measured."""
+    commit, dirty = git_state(source)
+    return {
+        "commit": commit,
+        "dirty": dirty,
+        "date": datetime.datetime.now().isoformat(timespec="seconds"),
+        "python": platform.python_version(),
+        "batch_size": BATCH_SIZE,
+        "compute_dtype": COMPUTE_DTYPE,
+        "warmup_steps": WARMUP_STEPS,
+        "baseline_steps": BASELINE_STEPS,
+    }

benchmarks/history/v1.2.0/torch_rtx3050.json

@@ -0,0 +1,22 @@
+{
+  "meta": {
+    "commit": "29c71e361d34f53efa34a84bf4551a9ac4d74ff7",
+    "dirty": true,
+    "date": "2026-11-01T21:48:18",
+    "python": "3.14.10",
+    "batch_size": 256,
+    "compute_dtype": "float16",
+    "warmup_steps": 20,
+    "baseline_steps": 128,
+    "gpu": "NVIDIA GeForce RTX 3050 6GB Laptop GPU",
+    "torch": "2.12.0+cu132",
+    "mode": "compiled",
+    "repeats": 3,
+    "warmup_s": 86.26663080000435
+  },
+  "baseline_ms": {
+    "median": 16.8658515625566,
+    "min": 16.82847109384511,
+    "max": 16.882614843780175
+  }
+}

benchmarks/history/v1.2.0/torch_rx9070xt.json

@@ -0,0 +1,22 @@
+{
+  "meta": {
+    "commit": "fc06f99038bca4e8c77283e548aeb97d323ccd1d",
+    "dirty": true,
+    "date": "2026-11-02T02:13:35",
+    "python": "3.14.4",
+    "batch_size": 256,
+    "compute_dtype": "float16",
+    "warmup_steps": 20,
+    "baseline_steps": 128,
+    "gpu": "AMD Radeon RX 9070 XT",
+    "torch": "2.12.0+rocm7.14.1",
+    "mode": "compiled",
+    "repeats": 3,
+    "warmup_s": 3.1427786790000027
+  },
+  "baseline_ms": {
+    "median": 7.6043785781223505,
+    "min": 7.583219632813609,
+    "max": 7.621538335932598
+  }
+}