collective
Browse / github.com/storytold/effectcraft / patch
Avoid quadratic Schur assembly for long camera tracks
camtrack: Schur reduction scales quadratically with track length; parallel path buys nothing ## Problem (long tracks) `camtrack/bundle.rs:350-363`: the Schur reduction visits every pair of W entries per point (`wv` is also re-allocated per point). For long tracks the explicit materialisation scales quadratically in the track length: 1000 points / 120,000 observations / 120 cameras: 22.32 G instructions, 2428 ms per three-iteration BA, 93 % in the Schur reduction (`adjust::{closure#5}`), 3.4 % Cholesky. A faster Cholesky does not help here (see the band issue). Related: `bundle.rs:196` (`parallel = n <= 700`) and `:264-293`, `:367-383`: each worker gets dense n^2 buffers, and the domin
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #2346, open)
- Base
991904888a20b7da72b2a4880b21c7d912154736- Size
- +336 −43 in 2 files
- Review
- pending
- Shared
- 2026-10-28
- Id
770fcee01f2c0a594e01c5e0
How it went
No other operator has reported using it yet.
Diff
raw · crates/track/src/camtrack/bundle.rs crates/track/src/camtrack/linalg.rs
@@ -174,8 +174,107 @@ fn cam_jac(l: &Layout, ci: usize, c: &BaCam, xc: V3, a: &[[f64; 3]; 2], p: [f64; } } +struct PtSys {+ v: M3,+ bp: V3,+ w: Vec<(usize, V3)>,+}++struct PreconditionerBlock {+ start: usize,+ dim: usize,+ inverse: Vec<f64>,+}++fn camera_preconditioner(u: &[f64], l: &Layout) -> Option<Vec<PreconditionerBlock>> {+ let mut ranges = Vec::new();+ let pose_dim = if l.rot_only { 3 } else { 6 };+ for (&pose, &focal) in l.start.iter().zip(&l.fidx) {+ if pose != usize::MAX {+ ranges.push((pose, pose_dim + usize::from(focal != usize::MAX)));+ } else if focal != usize::MAX {+ ranges.push((focal, 1));+ }+ }+ if let Some(start) = l.shared.or(l.dist) {+ ranges.push((start, l.n - start));+ }+ ranges+ .into_iter()+ .map(|(start, dim)| {+ let mut matrix = Vec::with_capacity(dim * dim);+ for i in 0..dim {+ for j in 0..dim {+ matrix.push(*u.get((start + i) * l.n + start + j)?);+ }+ }+ let mut inverse = vec![0.0; dim * dim];+ for col in 0..dim {+ let mut a = matrix.clone();+ let mut b = vec![0.0; dim];+ *b.get_mut(col)? = 1.0;+ if !cholesky_solve(&mut a, dim, &mut b) {+ return None;+ }+ for (row, value) in b.into_iter().enumerate() {+ inverse[row * dim + col] = value;+ }+ }+ Some(PreconditionerBlock { start, dim, inverse })+ })+ .collect()+}++/// Apply the point-eliminated system without materialising W V^-1 W^T.+fn iterative_schur(u: &[f64], rhs: &[f64], systems: &[PtSys], vinv: &[Option<M3>], l: &Layout) -> Option<Vec<f64>> {+ if l.n == 0 || l.n.checked_mul(l.n) != Some(u.len()) || rhs.len() != l.n || systems.len() != vinv.len() {+ return None;+ }+ let blocks = camera_preconditioner(u, l)?;+ // U contains camera blocks, global parameters and the focal prior. Store its nonzeros+ // once; all subsequent products are linear in these and the point-camera connections.+ let rows: Vec<Vec<(usize, f64)>> =+ u.chunks_exact(l.n).map(|row| row.iter().enumerate().filter(|(_, v)| **v != 0.0).map(|(i, v)| (i, *v)).collect()).collect();+ let mut reduced_rhs = rhs.to_vec();+ for (ps, vi) in systems.iter().zip(vinv) {+ let Some(vi) = vi else { continue };+ let vb = mv(vi, ps.bp);+ for (i, w) in &ps.w {+ *reduced_rhs.get_mut(*i)? -= dot(*w, vb);+ }+ }+ conjugate_gradient(+ &reduced_rhs,+ |x, y| {+ for (i, row) in rows.iter().enumerate() {+ y[i] = row.iter().map(|(j, v)| v * x[*j]).sum();+ }+ for (ps, vi) in systems.iter().zip(vinv) {+ let Some(vi) = vi else { continue };+ let q = ps.w.iter().fold([0.0; 3], |q, (i, w)| add(q, scale(*w, x[*i])));+ let q = mv(vi, q);+ for (i, w) in &ps.w {+ y[*i] -= dot(*w, q);+ }+ }+ },+ |r, z| {+ z.fill(0.0);+ for block in &blocks {+ for i in 0..block.dim {+ z[block.start + i] = (0..block.dim).map(|j| block.inverse[i * block.dim + j] * r[block.start + j]).sum();+ }+ }+ },+ )+}+ /// Run Levenberg–Marquardt. `cams` and `pts` are updated in place. pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> BaReport {+ adjust_impl(cams, pts, obs, o, true)+}++fn adjust_impl(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts, allow_iterative: bool) -> BaReport { let l = layout(cams.len(), o); let npts = pts.len(); // Observations by point.@@ -199,11 +298,6 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> iters += 1; // ---- Normal equations (Gauss–Newton with IRLS weights). // Per point: V (3×3), bp (3), W entries (cam-side index → 3).- struct PtSys {- v: [[f64; 3]; 3],- bp: V3,- w: Vec<(usize, V3)>,- } let build_point = |j: usize, s_acc: &mut Vec<f64>, bc: &mut Vec<f64>| -> PtSys { let mut ps = PtSys { v: [[0.0; 3]; 3], bp: [0.0; 3], w: vec![] }; let mut jc: Vec<(usize, [f64; 2])> = Vec::with_capacity(8);@@ -319,6 +413,10 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> let base_s = s_mat.clone(); let base_bc = bc.clone(); let mean_diag = ((0..n).map(|i| base_s[i * n + i]).sum::<f64>() / n.max(1) as f64).max(1e-12);+ let schur_pairs = systems.iter().fold(0usize, |sum, ps| sum.saturating_add(ps.w.len().saturating_mul(ps.w.len().saturating_add(1)) / 2));+ // Reserve the iterative solve for expensive long-track reductions; small systems+ // retain the direct path and its rounding. Failed solves use the direct path too.+ let mut use_iterative = allow_iterative && !o.fix_points && n >= 120 && schur_pairs > 4_000_000; // ---- Try steps with increasing damping. let mut accepted = false; for _try in 0..12 {@@ -346,50 +444,62 @@ pub fn adjust(cams: &mut [BaCam], pts: &mut [V3], obs: &[BaObs], o: &BaOpts) -> }) .collect() };- if !o.fix_points {- let reduce = |range: std::ops::Range<usize>, s: &mut [f64], rhs: &mut [f64]| {- for j in range {- let (Some(vi), ps) = (&vinv[j], &systems[j]) else { continue };- let vb = mv(vi, ps.bp);- let wv: Vec<V3> = ps.w.iter().map(|(_, w)| mv(vi, *w)).collect();- for (a, (ia, wa)) in ps.w.iter().enumerate() {- rhs[*ia] -= dot(*wa, vb);- for (ib, wb) in &ps.w[a..] {- let d = dot(wv[a], *wb);- s[ia * n + ib] -= d;- if ia != ib {- s[ib * n + ia] -= d;+ let iterative_step = if use_iterative { iterative_schur(&s, &rhs, &systems, &vinv, &l) } else { None };+ if dbg && use_iterative {+ eprintln!("ba iterative n={n} pairs={schur_pairs} solved={}", iterative_step.is_some());+ }+ if use_iterative && iterative_step.is_none() {+ use_iterative = false;+ }+ let dc = if let Some(dc) = iterative_step {+ dc+ } else {+ if !o.fix_points {+ let reduce = |range: std::ops::Range<usize>, s: &mut [f64], rhs: &mut [f64]| {+ for j in range {+ let (Some(vi), ps) = (&vinv[j], &systems[j]) else { continue };+ let vb = mv(vi, ps.bp);+ let wv: Vec<V3> = ps.w.iter().map(|(_, w)| mv(vi, *w)).collect();+ for (a, (ia, wa)) in ps.w.iter().enumerate() {+ rhs[*ia] -= dot(*wa, vb);+ for (ib, wb) in &ps.w[a..] {+ let d = dot(wv[a], *wb);+ s[ia * n + ib] -= d;+ if ia != ib {+ s[ib * n + ia] -= d;+ } } } }- }- };- let chunks = rayon::current_num_threads().clamp(1, 8);- if n <= 1200 && npts >= 64 && chunks > 1 {- let per = npts.div_ceil(chunks);- let parts: Vec<(Vec<f64>, Vec<f64>)> = (0..chunks)- .into_par_iter()- .map(|t| {- let (mut ds, mut dr) = (vec![0.0; n * n], vec![0.0; n]);- reduce(t * per..((t + 1) * per).min(npts), &mut ds, &mut dr);- (ds, dr)- })- .collect();- for (ds, dr) in parts {- s.par_iter_mut().zip(ds.par_iter()).for_each(|(x, y)| *x += y);- for (x, y) in rhs.iter_mut().zip(dr) {- *x += y;+ };+ let chunks = rayon::current_num_threads().clamp(1, 8);+ if n <= 1200 && npts >= 64 && chunks > 1 {+ let per = npts.div_ceil(chunks);+ let parts: Vec<(Vec<f64>, Vec<f64>)> = (0..chunks)+ .into_par_iter()+ .map(|t| {+ let (mut ds, mut dr) = (vec![0.0; n * n], vec![0.0; n]);+ reduce(t * per..((t + 1) * per).min(npts), &mut ds, &mut dr);+ (ds, dr)+ })+ .collect();+ for (ds, dr) in parts {+ s.par_iter_mut().zip(ds.par_iter()).for_each(|(x, y)| *x += y);+ for (x, y) in rhs.iter_mut().zip(dr) {+ *x += y;+ } }+ } else {+ reduce(0..npts, &mut s, &mut rhs); }- } else {- reduce(0..npts, &mut s, &mut rhs); }- }- let mut dc = rhs;- if n > 0 && !cholesky_solve(&mut s, n, &mut dc) {- lambda *= 10.0;- continue;- }+ let mut dc = rhs;+ if n > 0 && !cholesky_solve(&mut s, n, &mut dc) {+ lambda *= 10.0;+ continue;+ }+ dc+ }; // Apply. let mut nc: Vec<BaCam> = cams.to_vec(); for (i, c) in nc.iter_mut().enumerate() {@@ -464,6 +574,111 @@ mod tests { use super::*; use crate::camtrack::geometry::Rng; + #[test]+ fn long_track_adjustment_matches_direct_reprojections_and_outliers() {+ rayon::ThreadPoolBuilder::new().num_threads(1).build().unwrap().install(|| {+ for (focal, fix_centers, distortion) in+ [(FocalMode::Fixed, false, false), (FocalMode::Shared, false, true), (FocalMode::PerCamera(30.0), false, true), (FocalMode::Shared, true, true)]+ {+ let mut rng = Rng(71);+ let mut noise = || (rng.next_u64() % 100_000) as f64 / 100_000.0 - 0.5;+ let truth: Vec<V3> = (0..320).map(|_| [noise() * 8.0, noise() * 5.0, 10.0 + noise() * 5.0]).collect();+ let true_cams: Vec<_> = (0..60)+ .map(|i| BaCam {+ pose: Pose {+ r: rodrigues([0.001 * i as f64, -0.003 * i as f64, 0.0]),+ c: if fix_centers { [0.0; 3] } else { [0.03 * i as f64, 0.01 * (i as f64).sin(), 0.0] },+ },+ f: if matches!(focal, FocalMode::PerCamera(_)) { 900.0 * (1.0 + 0.05 * (i as f64 * 0.07).sin()) } else { 900.0 },+ dist: if distortion { Distortion { k1: 0.015, k2: -0.005, radius: 1100.0 } } else { Distortion::NONE },+ })+ .collect();+ let mut obs = Vec::new();+ for (ci, c) in true_cams.iter().enumerate() {+ for (pi, p) in truth.iter().enumerate() {+ let mut uv = residual(c, *p, [0.0; 2]).unwrap();+ uv[0] += noise() * 0.05;+ uv[1] += noise() * 0.05;+ if (ci * truth.len() + pi).is_multiple_of(37) {+ uv[0] += 25.0;+ uv[1] -= 30.0;+ }+ obs.push(BaObs { cam: ci as u32, pt: pi as u32, uv });+ }+ }+ let mut cams = true_cams.clone();+ for c in cams.iter_mut().skip(2) {+ c.pose.r = mm(&rodrigues([0.001, -0.001, 0.0005]), &c.pose.r);+ if !fix_centers {+ c.pose.c = add(c.pose.c, [0.002, -0.003, 0.001]);+ }+ }+ let mut pts: Vec<_> = truth.iter().map(|p| add(*p, [0.01, -0.01, 0.005])).collect();+ let (mut direct_cams, mut direct_pts) = (cams.clone(), pts.clone());+ let opts = BaOpts { focal, fix_centers, distortion, fixed: (0..60).map(|i| i < 2).collect(), max_iter: 10, ..Default::default() };+ let report = adjust(&mut cams, &mut pts, &obs, &opts);+ let direct = adjust_impl(&mut direct_cams, &mut direct_pts, &obs, &opts, false);+ assert!((report.cost - direct.cost).abs() < 1e-5 * direct.cost.max(1.0));+ for ob in &obs {+ let a = residual(&cams[ob.cam as usize], pts[ob.pt as usize], ob.uv).unwrap();+ let b = residual(&direct_cams[ob.cam as usize], direct_pts[ob.pt as usize], ob.uv).unwrap();+ assert!((a[0] - b[0]).hypot(a[1] - b[1]) < 0.001);+ assert_eq!(a[0].hypot(a[1]) <= 4.0, b[0].hypot(b[1]) <= 4.0);+ }+ for (a, b) in cams.iter().zip(direct_cams) {+ assert!(norm(sub(a.pose.c, b.pose.c)) < 0.001);+ assert!(norm(log_rot(&mm(&a.pose.r, &mt(&b.pose.r)))) < 1e-4);+ assert!((a.f - b.f).abs() < 1e-4 * b.f);+ assert!((a.dist.k1 - b.dist.k1).abs() < 1e-4 && (a.dist.k2 - b.dist.k2).abs() < 1e-4);+ }+ }+ });+ }++ #[test]+ fn matrix_free_schur_matches_explicit_system_and_rejects_bad_preconditioner() {+ let l = layout(20, &BaOpts::default());+ let n = l.n;+ let mut u = vec![0.0; n * n];+ for i in 0..n {+ u[i * n + i] = 10.0;+ }+ let systems: Vec<_> = (0..80)+ .map(|j| PtSys {+ v: I3,+ bp: [0.01, -0.02, 0.03],+ w: (0..n).map(|i| (i, [((j * 13 + i * 7) % 19) as f64 * 0.001 - 0.009, (i % 7) as f64 * 0.001, (j % 11) as f64 * 0.001])).collect(),+ })+ .collect();+ let vinv = vec![Some(I3); systems.len()];+ let rhs: Vec<_> = (0..n).map(|i| (i % 7) as f64 - 3.0).collect();+ let x = iterative_schur(&u, &rhs, &systems, &vinv, &l).unwrap();+ let mut explicit = u.clone();+ let mut reduced_rhs = rhs.clone();+ for ps in &systems {+ for (i, a) in &ps.w {+ reduced_rhs[*i] -= dot(*a, ps.bp);+ for (j, b) in &ps.w {+ explicit[i * n + j] -= dot(*a, *b);+ }+ }+ }+ let mut direct = reduced_rhs.clone();+ let mut factor = explicit.clone();+ assert!(cholesky_solve(&mut factor, n, &mut direct));+ for (a, b) in x.iter().zip(direct) {+ assert!((a - b).abs() < 1e-9);+ }+ let mut error2 = 0.0;+ for (i, row) in explicit.chunks_exact(n).enumerate() {+ let ax: f64 = row.iter().zip(&x).map(|(a, x)| a * x).sum();+ error2 += (ax - reduced_rhs[i]).powi(2);+ }+ assert!(error2 < 1e-18 * reduced_rhs.iter().map(|v| v * v).sum::<f64>());+ u[0] = -1.0;+ assert!(iterative_schur(&u, &rhs, &systems, &vinv, &l).is_none());+ }+ #[test] fn bundle_adjustment_converges_from_noise() { let mut rng = Rng(99);@@ -353,6 +353,60 @@ pub fn cholesky_solve(a: &mut [f64], n: usize, b: &mut [f64]) -> bool { b.iter().all(|v| v.is_finite()) } +/// Preconditioned conjugate gradients with an independently recomputed residual check.+/// Failure to converge (or non-finite/non-positive curvature) requests the direct fallback.+pub(crate) fn conjugate_gradient(rhs: &[f64], apply: impl Fn(&[f64], &mut [f64]), precondition: impl Fn(&[f64], &mut [f64])) -> Option<Vec<f64>> {+ let n = rhs.len();+ let norm2 = dot_n(rhs, rhs);+ if !norm2.is_finite() {+ return None;+ }+ if rhs.iter().all(|v| *v == 0.0) {+ return Some(vec![0.0; n]);+ }+ let tolerance2 = norm2 * 1e-22;+ if tolerance2 == 0.0 {+ return None;+ }+ let mut x = vec![0.0; n];+ let mut r = rhs.to_vec();+ let mut z = vec![0.0; n];+ let mut ap = vec![0.0; n];+ precondition(&r, &mut z);+ let mut p = z.clone();+ let mut rz = dot_n(&r, &z);+ for _ in 0..n.saturating_mul(2).min(512) {+ if !rz.is_finite() || rz <= 0.0 {+ return None;+ }+ apply(&p, &mut ap);+ let curvature = dot_n(&p, &ap);+ if !curvature.is_finite() || curvature <= 0.0 {+ return None;+ }+ let alpha = rz / curvature;+ for i in 0..n {+ x[i] += alpha * p[i];+ r[i] -= alpha * ap[i];+ }+ if dot_n(&r, &r) <= tolerance2 {+ apply(&x, &mut ap);+ for i in 0..n {+ ap[i] = rhs[i] - ap[i];+ }+ return (x.iter().all(|v| v.is_finite()) && dot_n(&ap, &ap) <= 100.0 * tolerance2).then_some(x);+ }+ precondition(&r, &mut z);+ let next_rz = dot_n(&r, &z);+ let beta = next_rz / rz;+ for i in 0..n {+ p[i] = z[i] + beta * p[i];+ }+ rz = next_rz;+ }+ None+}+ /// A similarity transform `x ↦ s R x + t`. #[derive(Clone, Copy, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub struct Similarity {@@ -428,6 +482,30 @@ pub fn from_euler_xyz(e: V3) -> M3 { mod tests { use super::*; + #[test]+ fn conjugate_gradient_checks_residual_and_bad_curvature() {+ let n = 64;+ let rhs: Vec<f64> = (0..n).map(|i| (i % 11) as f64 - 5.0).collect();+ let apply = |x: &[f64], y: &mut [f64]| {+ for i in 0..n {+ y[i] = 4.0 * x[i] + if i > 0 { x[i - 1] } else { 0.0 } + if i + 1 < n { x[i + 1] } else { 0.0 };+ }+ };+ let x = conjugate_gradient(&rhs, apply, |r, z| {+ for (r, z) in r.iter().zip(z) {+ *z = *r / 4.0;+ }+ })+ .unwrap();+ let mut product = vec![0.0; n];+ apply(&x, &mut product);+ let residual: f64 = rhs.iter().zip(product).map(|(b, ax)| (b - ax).powi(2)).sum();+ assert!(residual < 1e-18 * dot_n(&rhs, &rhs));+ assert!(conjugate_gradient(&[1.0], |x, y| y[0] = -x[0], |r, z| z.copy_from_slice(r)).is_none());+ assert!(conjugate_gradient(&[f64::NAN], |_, _| {}, |_, _| {}).is_none());+ assert_eq!(conjugate_gradient(&[0.0], |_, _| {}, |_, _| {}).unwrap(), vec![0.0]);+ }+ #[test] fn svd_reconstructs() { let a = [3.0, 1.0, 2.0, -1.0, 4.0, 0.5, 2.0, 2.0, 7.0, 0.1, -3.0, 1.0];