Profiling and Performance#
benchmarks/profile_pipeline.py profiles a full run_simulation call
rather than a micro-benchmark of one kernel, because the relevant costs
are distributional: which stage dominates depends on N, on how often
sticking fails, and on how many retries a regime forces. It rolls
cProfile output up into the pipeline’s own stages, so the report
identifies what to optimize rather than only what is hot.
devenv shell -- uv run python benchmarks/profile_pipeline.py --n 1024 --regime medium
Findings#
The first profiling pass found four separate problems, none of them in the overlap kernels (already JIT-compiled, and never the bottleneck).
np.cross on 3-vectors costs 21 µs#
np.cross supports 2- and 3-component inputs over arbitrary axes with
broadcasting, and pays for that generality on every call. Against three
lines of explicit component arithmetic:
Operation |
Time |
Ratio |
|---|---|---|
|
21.0 µs |
— |
explicit |
1.7 µs |
12.4x |
|
1.28 µs |
— |
|
0.63 µs |
2.0x |
The CCA sticking path made ~19,500 norm and ~1,550 cross calls per
N=512 aggregate, all on plain 3-vectors. geometry.norm3/cross3 now
serve those sites. These are deliberately not JIT-compiled: they are
called with plain Python floats and small arrays from interpreted code,
where numba’s dispatch overhead would exceed the gain.
The per-aggregate quality record was O(N²) in memory#
quality.max_self_overlap built an (N, N, 3) difference array, and
runs on every aggregate — a regression introduced together with the
quality record itself:
N |
Before |
After (pdist) |
|---|---|---|
512 |
12.2 ms |
~0.4 ms |
1024 |
49.3 ms |
~1.1 ms |
2048 |
211 ms, 101 MB temporary |
~5.5 ms |
scipy.spatial.distance.pdist walks the upper triangle directly in C.
Disabled telemetry was still being computed#
_leaf_mask_for_cluster is O(n²) per cluster and ran on every sticking
call, as did _candidate_score per candidate attempt. Under the
production defaults neither result is read: the baseline candidate
policy ignores leaf classification, and the leaf/score counters are
only printed when their profile_* flags are set. Both are now
computed only when something consumes them — an experimental candidate
policy, a profiling flag, or an attached merge log.
The sticking placement was scalar glue around fast kernels#
_cca_sticking_v1 ran 2,176 times per N=1024 aggregate and accounted
for ~26% of the run, almost none of it arithmetic: ~30k norm calls,
~24k np.array allocations and ~20k np.zeros, all on 3-vectors and
small temporaries between operations that were already fast.
cca_kernels.cca_sticking_v1_kernel fuses the whole default
(ext_case=0) placement — translate cluster 2, sample the first
contact point, rotate cluster 1, sample the second contact point,
rotate cluster 2 — into one compiled function, including both cluster
transforms as tight loops. Its self time dropped 5.8×, from 133 ms to
23 ms.
Randomness is hoisted out of the kernel: the two angles the
sphere-sphere intersections would have sampled are passed in as
arguments, so the caller retains sole ownership of the RNG stream and
seeded runs stay reproducible. This also makes the kernel a pure
function, which is what allows tests/test_sticking_kernel.py to pin
it against the interpreted reference exactly — seeding both identically
makes the comparison bit-for-bit rather than statistical. ext_case=1
(spherical-cap sampling) still runs the interpreted path, and
_cca_sticking_v1_interpreted is retained as the readable definition
of correct behaviour.
Measured effect#
End-to-end medians, 5 seeds per cell, single-threaded:
Regime |
N |
Before |
After |
Speedup |
|---|---|---|---|---|
easy |
256 |
130.9 ms |
111.4 ms |
1.18x |
easy |
512 |
267.4 ms |
203.5 ms |
1.31x |
easy |
1024 |
498.8 ms |
329.9 ms |
1.51x |
medium |
256 |
174.4 ms |
125.2 ms |
1.39x |
medium |
512 |
450.7 ms |
251.1 ms |
1.80x |
medium |
1024 |
998.7 ms |
571.9 ms |
1.75x |
The gain grows with N and with difficulty, as expected: the O(N²) fixes scale with size, and the sticking kernel scales with how many candidate placements a regime forces.
Results are unchanged: a seeded run remains bit-identical to itself and
produces the same overlap-free geometry. Aggregates generated by older
versions at the same seed may differ in the last bits, since
math.sqrt of an explicit sum and np.linalg.norm need not associate
identically; this shifts a trajectory without changing its statistics.
A negative result on log-call guards#
Recorded to prevent repetition. Three logger.trace(f"...") calls in
the rotation loop evaluate their f-string before trace() can check
the level — a genuine and common defect — but guarding them produced no
measurable change, because only ~900 rotations run per aggregate and a
float format costs ~1 µs. The guards were kept as a correctness fix,
not as an optimization.
The expectation that they would matter came from trusting cProfile’s
self time, which charges some per-call instrumentation overhead to
the calling function and made _perform_cca_sticking’s loop body
appear to hold ~47% of the run. Wall-clock A/B measurement is the
arbiter; the profiler is useful for pointing at candidates only.
Remaining time#
After the changes above, the profile at N=1024/medium is dominated by
_perform_cca_sticking’s own loop body, with the geometry and sticking
stages largely gone (_cca_sticking_v1 fell from 14% of the run to
under 4%). The self-time caveat above applies when reading that
figure: a loop making thousands of calls is exactly the shape cProfile
over-attributes.
Further gains would require compiling the whole candidate/rotation loop, not just the placement inside it. That is a substantially harder change — the loop interleaves RNG draws, overlap checks and the retry policy, and would have to keep the RNG draw order identical to remain reproducible — and the payoff has shrunk: PCA no longer registers on the profile, and a representative N=1024 aggregate builds in about half a second.