Schur elimination + benchmarking#57
Merged
Merged
Conversation
Automatic Schur-complement elimination of dominant block-diagonal variable types (e.g. landmarks in bundle adjustment): solve() forms a reduced "kept" system, solved dense or matrix-free CG, then back-substitutes. Supports multiple kept and multiple eliminated types; falls back to the full system for cholmod or when no dominant block exists. Performance: small-contraction batched products are written as broadcast-multiply-sums (utils._batched_gram/_outer_last/_matmul) rather than einsum — XLA lowers tiny-contraction batched GEMMs to ~15-30x slower kernels on GPU. Applied to Schur assembly and the block-Jacobi preconditioner. ~4x on dense-S assembly. Levenberg-Marquardt fixes in the shared loop (improve every solver): - Gate the cost AND parameter convergence criteria on step acceptance; a rejected proposal's deltas said nothing about progress and could freeze LM into re-proposing the same rejected step. - Numerically stable predicted reduction (2 dx^T ATb - |A dx|^2), avoiding the catastrophic cancellation of differencing two O(cost) sums; also fixes a double-application of the column scaler in the old form. - Nielsen-style accelerating lambda escalation (factor doubles per consecutive rejection; no-op when the first trial is accepted). - Don't end an augmented-Lagrangian solve at lambda_max. - Compute the Jacobian column scaler once (it's frozen at the initial linearization) instead of every iteration. Robustness: float32 Schur S uses |diag| for Jacobi scaling (a negative diagonal from float32 cancellation no longer blows up a row), and V is inverted with a tiny damping floor so Gauss-Newton (lambda=0) on an under-observed landmark can't divide by a singular block. 53/53 tests pass (exactness vs full dense ~1e-10). _py310 transpiled.
Instrumentation (the suite's measurement primitive): - SolveSummary.time_history: per-outer-iteration host wall-clock timestamps via jax.pure_callback, so one solve yields a cost-vs-time trace (no matched-k re-solving). Opt-in via TerminationConfig.record_time_history (default off — the callback is a host sync point unwanted on the normal solve path and awkward under vmap/pmap). The callback is anchored on the iterate's cost so XLA can't reorder/CSE/hoist it ahead of the step's work. Benchmark suite (benchmarks/suite/, one tyro CLI; see benchmarks/README.md): - python -m benchmarks.suite [--quick] [--gate] [--update-baseline] [--only ...] - Workloads reduce a jaxls run to named scalar Metrics; baseline.json is the committed regression reference; --gate exits nonzero on regression. - Default gate set = bundle_adjustment + float32_robustness (fast, in-process-reliable). pyroki_ik and example_notebooks are opt-in (slow and/or unreliable as co-tenant CUDA subprocesses). - Subprocess workloads share one _run_marker_subprocess helper that runs under _subprocess_env (free-GPU pinning + prealloc/command-buffer disable so children don't OOM against the CUDA-warmed parent); float32 reuses benchmarks/float32_check.run rather than re-spelling the solve. Study harnesses: device_sweep.py (CPU/GPU BA matrix + plots), trace_examples.py (per-example cost-vs-time, main vs PR), analytic_jac_ba.py (analytic vs autodiff Jacobian), make_regression_report.py. results.md and regression.md carry the CPU/GPU study, the LM/scaler investigation, the performance-ceiling (host-dispatch-bound) writeup, and main-vs-PR tables. Builds on the schur-elimination PR. 53/53 tests pass; _py310 transpiled.
Re-ran transpile_py310.py through the project's pinned ruff (uv run) so the generated _py310 files match the canonical formatter version, removing the whitespace churn from an earlier run under a different ruff.
Per-iteration timestamps were stored in a jax.Array on SolveSummary, set via a host callback during the solve. Without jax_enable_x64 that array is float32, and perf_counter() values (~5e5) have zero float32 resolution at millisecond scale, so fast solves recorded all-equal timestamps and their cost-vs-time plots came out blank (black_litterman, cvar_optimization). Move recording off the jitted path entirely: a `record_iteration_times()` context manager exposes a host-side Python list filled via jax.debug.callback (always float64, no SolveSummary field, zero overhead and no dtype dependence when not recording). The callback's presence is fixed at trace time and isn't part of the JIT cache key, so the context manager clears the solver cache on entry/exit to force a callback-bearing recompile inside the block and restore the zero-overhead executable after. Drops the SolveSummary.time_history field and TerminationConfig .record_time_history flag. Updates trace_examples.py (feature-detects the old main-side API) and analytic_jac_ba.py to the new API, and fixes a length-mismatch when a solve runs the full max_iterations (cost_history caps at max_iterations but the recorder does not). Regenerates the example-trace and analytic-Jacobian plots with real float64 wall-clock.
Convert every benchmark invocation in docstrings, DEV_NOTES, results.md and the suite CLI help to `uv run --extra dev --extra docs python ...`, matching how the project is actually run. No behavior change.
Surface the numbers a reader cares about directly on each plot instead of making them eyeball the curves: - BA comparison (ba_comparison_gpu.png): per-method speedup in the legend, measured as time to reach the baseline full-CG converged cost (4× / 196× / 624× for Schur+dense; 3× / 29× for Schur+CG), plus a dotted line at that baseline cost. Larger fonts, gridlines. - Example traces: PR-vs-main speedup per subplot title (time to reach main's final cost) and step counts in the legend. - Analytic-vs-autodiff Jacobian: per-solve ms in the legend and an honest verdict in the title — these are within noise (jaxls is host-dispatch-bound, so saving Jacobian device-FLOPs barely moves the wall-clock), reported as "matched within noise" rather than an inflated speedup. Larger figures/fonts and gridlines throughout.
Replace the vague "624× faster" legend tags with a dashed line at the baseline full-CG converged cost and a labeled dot where each method crosses it, showing the absolute wall-clock time (e.g. Trafalgar-138: 50 ms for Schur+dense, 1.1 s for Schur+CG, 31.3 s for full CG). The reader reads the times directly and can judge the gap themselves.
The threshold line was main full-CG's converged cost, which a slower- converging method might never reach (Schur+CG on Ladybug-49 had no labeled crossing). Set it to the highest of the three methods' converged costs — the best cost every method actually attains — so all three get a labeled crossing time. Relabel it "shared cost target".
CI's `uvx ruff` and pyright were failing:
- Bump pinned ruff 0.14.6 -> 0.15.17 (matches CI's uvx ruff) and reformat;
re-transpile the _py310 files with it (0.15 adds a blank line after
signatures whose docstring the transpiler strips), so check_transpile and
ruff format --check pass.
- pyright (run on the whole repo) flagged the new benchmark scripts. Fix all:
- _solvers.py: route the @jdc.jit .clear_cache() through a small helper
with a single type: ignore.
- workloads.py / baseline.py: assert the non-None branch after the skip /
regressed guards (the two correlated values pyright can't narrow).
- bal.py / analytic_jac_ba.py: type: ignore on **kwargs splats that only
ever carry known keys; coalesce as_text() | "".
- device_sweep.py: jax.Device (a nanobind type) isn't a valid type
expression to pyright — ignore on the annotation.
- __main__.py: the `only` field "docstring" was string + join (a discarded
expression that silently dropped the help text); make it a real
docstring.
- pyroki_ik.py / workloads.py: type: ignore on the optional pyroki import.
- test_schur.py: annotate the preconditioner loop tuple with its Literal
type.
The PR shipped five exploration scripts and ~2 MB of generated artifacts that the standing suite doesn't use. Remove them to keep the PR focused on the library change + the regression suite. Scripts removed (none imported by benchmarks/suite/): matched_iters.py and plot_frontier.py (the matched-k study, superseded by device_sweep.py, which runs the same methodology on CPU+GPU), profile_schur.py, analytic_jac_ba.py, and make_regression_report.py. Artifacts removed (regenerable; not read by the surviving code at runtime): example_traces.json, matched_*.json, results/regression/*.json, and all PNGs except the two headline plots (ba_comparison_gpu, example_traces) — ~1.5 MB. Kept: suite/baseline.json (the gate baseline), ba_initial_cost.json and device_*_tuned_*.json (read by device_sweep --replot). Also drops regression.md (generated by the removed make_regression_report.py from the removed regression/*.json) and updates the docstrings/README/results.md references that pointed at the deleted scripts.
Variable elimination now supports a sparse-direct (CHOLMOD) factorization of the reduced system, the Ceres/g2o-style "Schur + sparse-direct" combination: the block-diagonal landmark elimination runs on-device, and only the small camera system is handed to CHOLMOD. Enabled for linear_solver="cholmod" (previously elimination was disabled for it). The reduced system S is assembled as a symmetric COO matrix whose nonzero structure is precomputed on the host (build_sparse_s_pattern), with a single source of truth (_iter_S_block_sources) shared by the index builder and the per-iteration value assembler so they cannot drift. Symbolic factorization is cached and reused across iterations. benchmarks/device_sweep.py: add the cholmod-vs-Schur+cholmod study (run_cholmod_study, --cholmod) and redesign plot_study as a two-panel suboptimality-gap + speedup-to-solution figure. Tests: sparse-S equals dense-S, single-step exactness vs the full solve, and end-to-end convergence matching the dense Schur path. Full suite green (56).
- device_sweep.py: cast np.arange index to float for ax.text (pyright reportArgumentType). - Regenerate src/jaxls/_py310 with the project's uv-pinned ruff (0.15.17) so the transpiled sources match what the check_transpile CI job produces; the previous copy was generated with a newer ruff and differed by blank lines.
analyze(schur_elimination=...) now accepts:
- "auto" (default): infer a dominant block-diagonal type to eliminate;
- "off": solve the full system;
- a tuple of variable types, e.g. (LandmarkVar,): eliminate exactly those,
validated by build_elimination_plan (ValueError if a named type isn't
block-diagonal).
Clean break from the prior bool flag. The docstring notes that only a single
level of elimination is supported (no nested/multi-level Schur); the tuple
form is the explicit override.
The benchmarks drop the bool too: bal.load_bal/make_toy_ba and
device_sweep's loaders thread "auto"/"off" directly, and bal._analyze just
forwards the string when the installed jaxls has the parameter (jaxls@main
predates it). Tests cover the explicit-tuple, "off", and invalid-argument
paths.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added a
schur_elimination: Literal["auto", "off"] | tuple[type[Var], ...]argument toLeastSquaresProblem.analyze(). It's set to"auto"by default.This should significantly speed up problems where Hessians are partially block-diagonal. For example, bundle adjustment:
The PR also makes some small numerical changes that we discovered along the way. Before/after of example problems (as regression test):