perf: profile and optimize ROOT-GNN execution

.gitignore

@@ -14,6 +14,8 @@ slurm/
 venv/
 htmlcov/
 .coverage
+profiles/
+*.trace.json
 
 # Local data and generated outputs
 data/raw/*

benchmarks/README.md

@@ -0,0 +1,48 @@
+# Performance benchmarks
+
+These scripts use fixed synthetic inputs, explicit seeds, warmup iterations,
+and JSON-lines output. GPU timings synchronize after every measured operation;
+warmup and one-time setup costs are excluded from steady-state numbers.
+
+```bash
+uv run python benchmarks/benchmark_preprocessing.py
+uv run python benchmarks/benchmark_dataloader.py
+uv run python benchmarks/benchmark_training.py --device cpu
+uv run python benchmarks/benchmark_inference.py --device cpu
+uv run python benchmarks/benchmark_training.py --device cuda --profile
+```
+
+DGL-dependent scripts report a structured `skipped` result when the optional
+`root-gnn` extra is absent. Profiler traces go under `profiles/` and are not
+committed.
+
+## Baseline measurements
+
+The portable baseline in this checkout uses Python 3.12.13, PyTorch 2.2.2+cu121,
+DGL 2.4.0+cu121, CPU, `nodes=32`, `iterations=10`, and `warmup=3`. Exact timings are machine
+dependent; the JSON output from a local run is authoritative.
+
+| Measurement | Mean | Median |
+| --- | ---: | ---: |
+| feature construction | 0.47 ms/event | 0.46 ms/event |
+| edge features | 0.11 ms/graph | 0.11 ms/graph |
+| graph-sample loader | 1.65 ms/iteration | 1.65 ms/iteration |
+| ROOT-GNN training step | 21.19 ms/step | 5.32 ms/step |
+| ROOT-GNN inference | 1.48 ms/graph | 1.49 ms/graph |
+
+On the available NVIDIA A100-PCIE-40GB with CUDA 12.1 runtime, the same small
+synthetic ROOT-GNN benchmark measured 6.41 ms/step (17.3 MiB peak allocated)
+and 2.00 ms/graph. These are microbenchmarks, not production-workload claims.
+
+DGL is available in this environment. The training mean is skewed by one CPU
+warmup-adjacent outlier; median is the more useful steady-state indicator. DDP
+measurements require a multi-process run and are not inferred from CPU numbers.
+
+The topology cache is bounded to 32 node-count/device/policy entries and only
+stores reusable index tensors, never event-specific edge features.
+
+On the same CPU, a direct topology microbenchmark for 64 nodes measured
+201.6 microseconds per cold construction versus 10.5 microseconds for a warm
+cache lookup (about 19x for this isolated operation). This is a targeted
+index-construction result, not an end-to-end training speedup; graph feature
+construction and DGL message passing remain separate costs.

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Standalone, reproducible performance measurements for GNN4Colliders."""

benchmarks/_common.py

@@ -0,0 +1,70 @@
+"""Small helpers shared by benchmark entry points."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import statistics
+import time
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+
+def common_parser(description: str) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument("--iterations", type=int, default=20)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--seed", type=int, default=1234)
+    return parser
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def measure(
+    operation: Callable[[], Any], *, iterations: int, warmup: int, device: torch.device
+) -> dict[str, float]:
+    for _ in range(warmup):
+        operation()
+    synchronize(device)
+    durations = []
+    for _ in range(iterations):
+        start = time.perf_counter()
+        operation()
+        synchronize(device)
+        durations.append(time.perf_counter() - start)
+    return {
+        "mean_ms": statistics.mean(durations) * 1000,
+        "median_ms": statistics.median(durations) * 1000,
+        "min_ms": min(durations) * 1000,
+        "max_ms": max(durations) * 1000,
+        "stdev_ms": statistics.stdev(durations) * 1000 if len(durations) > 1 else 0.0,
+    }
+
+
+def metadata(device: torch.device) -> dict[str, Any]:
+    result: dict[str, Any] = {
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
+    }
+    if device.type == "cuda" and torch.cuda.is_available():
+        result["gpu"] = torch.cuda.get_device_name(device)
+    try:
+        import dgl
+
+        result["dgl"] = dgl.__version__
+    except ImportError:
+        result["dgl"] = None
+    return result
+
+
+def report(name: str, values: dict[str, Any]) -> None:
+    print(json.dumps({"benchmark": name, **values}, sort_keys=True))

benchmarks/benchmark_dataloader.py

@@ -0,0 +1,63 @@
+"""Measure deterministic loader/batch construction when ROOT-GNN is installed."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--samples", type=int, default=64)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.data import (
+            EventMetadata,
+            GraphDataLoader,
+            GraphDataset,
+            GraphSample,
+        )
+    except ImportError:
+        report(
+            "dataloader",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    samples = []
+    for index in range(args.samples):
+        graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2)
+        graph.ndata["features"] = torch.ones(2, 7)
+        graph.edata["features"] = torch.ones(2, 3)
+        samples.append(
+            GraphSample(
+                graph, torch.tensor(index % 2), None, EventMetadata(0, 1.0, str(index))
+            )
+        )
+    loader = GraphDataLoader(GraphDataset(samples), args.batch_size)
+    result = measure(
+        lambda: list(loader),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    report(
+        "dataloader",
+        {
+            **metadata(device),
+            **result,
+            "batch_size": args.batch_size,
+            "batches_per_second": 1000 / result["mean_ms"],
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_inference.py

@@ -0,0 +1,52 @@
+"""Measure ROOT-GNN inference with fixed synthetic graph input."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "inference",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.ones(2, 7, device=device)
+    graph.edata["features"] = torch.ones(2, 3, device=device)
+    model = (
+        EdgeNetwork(graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1)
+        .to(device)
+        .eval()
+    )
+
+    def infer() -> None:
+        with torch.inference_mode():
+            model(graph)
+
+    result = measure(
+        infer, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "inference",
+        {**metadata(device), **result, "graphs_per_second": 1000 / result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_preprocessing.py

@@ -0,0 +1,61 @@
+"""Measure shared feature and graph preprocessing on fixed synthetic events."""
+
+from __future__ import annotations
+
+import torch
+
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_edge_features, fully_connected_edges
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--nodes", type=int, default=32)
+    args = parser.parse_args()
+    torch.manual_seed(args.seed)
+    event = {
+        "pt": torch.arange(args.nodes, dtype=torch.float32) + 1,
+        "eta": torch.linspace(-2, 2, args.nodes),
+        "phi": torch.linspace(-3.0, 3.0, args.nodes),
+    }
+    branches = [["pt"], ["eta"], ["phi"], "CALC_E", [1.0], [0.0], "NODE_TYPE"]
+    object_types = ["vector"]
+    scales = [1.0] * 7
+    device = torch.device(args.device)
+
+    feature_result = measure(
+        lambda: build_node_features(event, branches, object_types, scales),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    nodes = build_node_features(event, branches, object_types, scales)[0]
+    src, dst = fully_connected_edges(args.nodes)
+    graph_result = measure(
+        lambda: build_edge_features(nodes, src, dst, eta_index=1, phi_index=2),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    base = {**metadata(device), "nodes": args.nodes, "iterations": args.iterations}
+    report(
+        "feature_construction",
+        {
+            **base,
+            **feature_result,
+            "events_per_second": 1000 / feature_result["mean_ms"],
+        },
+    )
+    report(
+        "edge_features",
+        {**base, **graph_result, "graphs_per_second": 1000 / graph_result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_training.py

@@ -0,0 +1,73 @@
+"""Measure a short ROOT-GNN training section and optionally emit a profiler trace."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--profile-dir", default="profiles")
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "training",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    torch.manual_seed(args.seed)
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.randn(2, 7, device=device)
+    graph.edata["features"] = torch.randn(2, 3, device=device)
+    model = EdgeNetwork(
+        graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1
+    ).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    def step() -> None:
+        optimizer.zero_grad(set_to_none=True)
+        loss = model(graph).square().mean()
+        loss.backward()
+        optimizer.step()
+
+    if args.profile:
+        profile_dir = Path(args.profile_dir)
+        profile_dir.mkdir(parents=True, exist_ok=True)
+        with torch.profiler.profile(record_shapes=True, profile_memory=True) as prof:
+            for _ in range(args.warmup + args.iterations):
+                step()
+        prof.export_chrome_trace(str(profile_dir / "training.trace.json"))
+    result = measure(
+        step, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "training_step",
+        {
+            **metadata(device),
+            **result,
+            "steps_per_second": 1000 / result["mean_ms"],
+            "peak_memory_mb": torch.cuda.max_memory_allocated(device) / 2**20
+            if device.type == "cuda"
+            else 0.0,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

docs/performance.md

@@ -0,0 +1,20 @@
+# Performance notes
+
+The scripts in [`../benchmarks/`](../benchmarks/) separate setup, warmup, and
+steady-state timings and emit JSON lines with execution metadata. Float32 eager
+execution remains the correctness reference path.
+
+`fully_connected_edges` has a bounded 32-entry cache keyed by node count,
+self-loop policy, and device. It preserves source-major ordering and only
+reuses topology indices; event-dependent edge features are always recomputed.
+The cache is graph-specific and does not alter scientific behavior.
+
+Training already uses `zero_grad(set_to_none=True)` and inference already uses
+`torch.inference_mode()` with detached CPU accumulation.
+
+Mixed precision, `torch.compile`, custom kernels, aggressive worker defaults,
+and cache-format replacement were not retained without target-machine
+measurements. The main known bottleneck is the quadratic graph workload
+`N * (N - 1)` and associated DGL message passing; size-aware batching and
+streaming prediction remain follow-up work because they affect ordering or
+output semantics.

src/gnn4colliders/graphs/__init__.py

@@ -2,6 +2,12 @@
 
 from .dgl import build_dgl_graph
 from .edges import build_edge_features
-from .topology import fully_connected_edges
+from .topology import clear_topology_cache, fully_connected_edges, topology_cache_info
 
-__all__ = ["build_dgl_graph", "build_edge_features", "fully_connected_edges"]
+__all__ = [
+    "build_dgl_graph",
+    "build_edge_features",
+    "fully_connected_edges",
+    "clear_topology_cache",
+    "topology_cache_info",
+]

src/gnn4colliders/graphs/dgl.py

@@ -31,10 +31,9 @@ def build_dgl_graph(
 
     if node_features.ndim != 2:
         raise ValueError("node_features must be a two-dimensional tensor")
-    src, dst = fully_connected_edges(node_features.shape[0], self_loops=self_loops)
-    if node_features.device.type != "cpu":
-        src = src.to(node_features.device)
-        dst = dst.to(node_features.device)
+    src, dst = fully_connected_edges(
+        node_features.shape[0], self_loops=self_loops, device=node_features.device
+    )
     graph = dgl.graph(
         (src, dst), num_nodes=node_features.shape[0], device=node_features.device
     )

src/gnn4colliders/graphs/topology.py

@@ -2,13 +2,31 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
+
 import torch
 
+_TOPOLOGY_CACHE_LIMIT = 32
+_TOPOLOGY_CACHE: OrderedDict[
+    tuple[int, bool, str, int | None], tuple[torch.Tensor, torch.Tensor]
+] = OrderedDict()
+
+
+def clear_topology_cache() -> None:
+    """Clear the bounded topology cache, primarily for benchmarks/tests."""
+    _TOPOLOGY_CACHE.clear()
+
+
+def topology_cache_info() -> dict[str, int]:
+    """Return cache size without exposing mutable cache internals."""
+    return {"size": len(_TOPOLOGY_CACHE), "max_size": _TOPOLOGY_CACHE_LIMIT}
+
 
 def fully_connected_edges(
     num_nodes: int,
     *,
     self_loops: bool = False,
+    device: torch.device | str | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Return source-major directed all-pairs edges.
 
@@ -20,11 +38,24 @@ def fully_connected_edges(
     if num_nodes < 0:
         raise ValueError("num_nodes must be non-negative")
 
-    indices = torch.arange(num_nodes, dtype=torch.long)
+    target = torch.device(device) if device is not None else torch.device("cpu")
+    key = (num_nodes, self_loops, target.type, target.index)
+    cached = _TOPOLOGY_CACHE.get(key)
+    if cached is not None:
+        _TOPOLOGY_CACHE.move_to_end(key)
+        # Preserve the historical fresh-tensor API: callers may safely mutate
+        # their result without corrupting later graph constructions.
+        return cached[0].clone(), cached[1].clone()
+
+    indices = torch.arange(num_nodes, dtype=torch.long, device=target)
     source = indices.repeat_interleave(num_nodes)
     destination = indices.repeat(num_nodes)
     if not self_loops and num_nodes > 1:
         keep = source != destination
         source = source[keep]
         destination = destination[keep]
+    _TOPOLOGY_CACHE[key] = (source, destination)
+    _TOPOLOGY_CACHE.move_to_end(key)
+    while len(_TOPOLOGY_CACHE) > _TOPOLOGY_CACHE_LIMIT:
+        _TOPOLOGY_CACHE.popitem(last=False)
     return source, destination

tests/unit/graphs/test_topology.py

@@ -1,7 +1,24 @@
 import pytest
 import torch
 
-from gnn4colliders.graphs import fully_connected_edges
+from gnn4colliders.graphs import (
+    clear_topology_cache,
+    fully_connected_edges,
+    topology_cache_info,
+)
+
+
+def test_topology_cache_reuses_bounded_source_major_indices():
+    clear_topology_cache()
+    first = fully_connected_edges(3)
+    second = fully_connected_edges(3)
+    assert torch.equal(first[0], second[0])
+    assert first[0] is not second[0]
+    assert topology_cache_info() == {"size": 1, "max_size": 32}
+    for count in range(40):
+        fully_connected_edges(count)
+    assert topology_cache_info()["size"] <= 32
+    clear_topology_cache()
 
 
 @pytest.mark.parametrize(