Skip to content

perf: rework the carré du champ inner loop; retire the porting plan - #26

Merged
wardjm merged 2 commits into
mainfrom
perf-cdc-tune-kernel
Jul 26, 2026
Merged

wardjm merged 2 commits into
mainfrom
perf-cdc-tune-kernel

Conversation

@wardjm

@wardjm wardjm commented Jul 26, 2026

Copy link
Copy Markdown
Owner

Two independent changes: dropping PORTING_PLAN.md now that the port is done, and a measured performance pass over the two places the time actually goes.

Retiring the porting plan

The port is complete and the repo diverges from upstream from here, so the phased plan has nothing left to track. Scope lives in README.md; the places where the Julia deliberately does not match the Python live in docs/src/upstream-bugs.md, which is where the one substantive cross-reference (the fixture-correction guard) now points.

Where the time goes

Profiling a 3000–5000 point cloud found the hot spots are not the einsums:

  • from_point_cloud is ~50% Arpack eigs and ~44% markov_chain — almost all of the latter is tune_kernel's exp sweep (80 ε × 60k entries, run twice).
  • Every differential operator is dominated by carre_du_champ_knn, not the OMEinsum contractions. For lie_bracket it was ~93% of the build.

carre_du_champ_knn

The per-point work now runs in a point-last layout. The neighbour gather becomes a contiguous column copy instead of a strided walk over an (n, F) array, and the rank-k update accumulates into a dense H×F buffer rather than the strided cdc[p, :, :] view — which had been forcing LinearAlgebra's generic matmul fallback once per point. Only that buffer is scattered back, and the 1/2ρ scaling moves to a single vectorised pass over the contiguous axis.

Two thresholds in here fell out of measurement rather than taste, and both are commented as such: BLAS beats a hand-rolled rank-1 accumulation only above ~2^11 flops (below that the per-point mul! call overhead dominates), and the mean-centring wants opposite loop orders for wide and narrow tensors.

Threading

The cdc point loop and the ε sweep are both embarrassingly parallel, so both are threaded. The cdc split is gated on F*H*k — ungated, spawning for small tensors cost more than it saved. Neither split changes results.

Also included: skipping exp terms that underflow to exactly zero (bit-identical, worth only a few percent), and hoisting a squared-distance matrix markov_chain was building twice.

Results

n=4000, cold geometry per rep, against a stashed baseline in the same session:

1 thread vs base 8 threads
grad 0.0072 −30% 0.0025
d(1) 0.0131 −20% 0.0091
laplacian(1) 0.0395 −42% 0.0241
hessian 0.1834 −10% 0.1597
levi_civita 0.0282 −8% 0.0170
lie_bracket 0.0427 −35% 0.0203
build 0.2545 ~flat 0.1943

Testing

Full suite green at both 1 and 8 threads — 2473 tests, parity fixtures and doctests included. That is the real gate on the muladd and loop-reordering changes, since they perturb floating-point association.

Notes for review

  • Type stability is not a lever here. The numeric kernels all infer concretely, so the usual performance-tips checklist finds nothing. The abstract fields that do exist (LinearOperator._weak::Union{Nothing,AbstractMatrix}, GammaCache::Any) are per-operator caches hit once per big matrix; tightening them buys nothing measurable, so they are untouched.
  • The build is now floored by Arpack, which is external and single-threaded. Replacing it (ArnoldiMethod/KrylovKit) is the next lever if build time matters more than operator time, but it carries real reproducibility risk — the current eigs call is deliberately pinned to a fixed v0 and a tight tolerance, and the comment there explains why. Deliberately out of scope.

wardjm added 2 commits July 26, 2026 13:22
The port is complete and the repo now diverges from upstream, so the phased
plan has nothing left to track. Scope lives in README.md; the places where the
Julia deliberately does not match the Python live in docs/src/upstream-bugs.md,
which is where the one substantive cross-reference (the fixture-correction
guard) now points.
Profiling a 3000-5000 point cloud puts the time in two places, neither of them
where the einsums are: from_point_cloud is ~50% Arpack and ~44% markov_chain
(almost all of that tune_kernel's exp sweep), and every differential operator is
dominated by carre_du_champ_knn -- ~93% of the lie_bracket build.

carre_du_champ_knn now does its per-point work in a point-last layout. The
neighbour gather becomes a contiguous column copy instead of a strided walk over
an (n, F) array, and the rank-k update accumulates into a dense H x F buffer
rather than the strided cdc[p, :, :] view, which had been forcing LinearAlgebra's
generic matmul fallback once per point. Only that buffer is scattered back, and
the 1/2rho scaling moves to a single vectorised pass over the contiguous axis.
Two thresholds fell out of measurement rather than taste: BLAS beats a hand
rolled rank-1 accumulation only above ~2^11 flops, and the mean-centring wants
opposite loop orders for wide and narrow tensors.

The point loop and the epsilon sweep are both embarrassingly parallel, so both
are threaded -- the cdc one gated on F*H*k so small tensors don't pay for a
spawn, which otherwise cost more than it saved. Neither split changes results.

Also skips exp terms that underflow to exactly zero (bit-identical, worth only a
few percent) and hoists a squared-distance matrix markov_chain built twice.

At n=4000, single-threaded, against cold geometry: grad -30%, d(1) -20%,
laplacian(1) -42%, hessian -10%, levi_civita -8%, lie_bracket -35%. With 8
threads lie_bracket is 0.020s against a 0.065s baseline. Full suite green at
both 1 and 8 threads, parity fixtures and doctests included.
@wardjm
wardjm merged commit 1c3cb44 into main Jul 26, 2026
3 checks passed
@wardjm
wardjm deleted the perf-cdc-tune-kernel branch August 3, 2026 20:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant