Skip to content

Add fused BA path#85

Merged
heiwang1997 merged 7 commits into
mainfrom
code/fused-ba
May 22, 2026
Merged

Add fused BA path#85
heiwang1997 merged 7 commits into
mainfrom
code/fused-ba

Conversation

@heiwang1997

@heiwang1997 heiwang1997 commented May 20, 2026

Copy link
Copy Markdown
Collaborator

Summary:

  • Add a fused CUDA BA path exposed by pipeline.slam.ba.fused.
  • Preserve the existing generic PyTorch BA solver as the default and fall back when the graph layout is unsupported.
  • Add slam_ext.ba_extended for optional depth masking and intrinsic optimization.
  • Add parity tests comparing fused BA against the generic solver on real-run-shaped graph-buffer data, including motion-only, limited disparity, intrinsic optimization, sparse-track combination, and a zero-weight self-edge corner case.

Notes:

  • The fused path keeps the DROID kernel's internal 0.001 flow-weight convention guarded in Python.
  • The CUDA path now includes dense-disparity damping alongside sensor-depth regularization to match ViPE's generic solver semantics.
  • Verbose fused BA logging now reports energy computed from the CUDA path rather than launching an extra Python-side energy recomputation.

Benchmark:

  • Video: /home/huangjh/Downloads/lyra2-temple.mp4 with pipeline=default (762 frames processed).
  • Common overrides: pipeline.output.save_viz=false, pipeline.post.depth_align_model=null, pipeline.output.save_artifacts=false, pipeline.output.save_slam_map=false; SLAM keyframe depth remains at the pipeline=default setting.
  • Command env: CUDA_HOME=/usr/local/cuda-12.8 VIPE_EXT_JIT=1 MAX_JOBS=8.
  • Extension JIT was already built (ninja: no work to do), so compile time is not included.

End-to-end wall time:

config run 1 run 2 mean
pipeline.slam.ba.fused=false 115.74s 116.03s 115.88s
pipeline.slam.ba.fused=true 78.60s 78.20s 78.40s

End-to-end delta: 37.48s saved, 32.3% lower wall time, 1.48x speedup.

Isolated BA timing:

run ba.fused=false BA time ba.fused=true BA time saved BA speedup
1 45.05s 8.01s 37.04s 5.63x
2 45.39s 7.75s 37.64s 5.86x
mean 45.22s 7.88s 37.34s 5.74x

BA-only method: timed GraphBuffer.bundle_adjustment with CUDA synchronization before and after each call, using the same video and common overrides. Each run made 1723 BA calls. Isolated BA time drops by about 82.6%, and the BA-only saving accounts for essentially all of the measured end-to-end saving on this benchmark.

Validation:

  • CUDA_HOME=/usr/local/cuda-12.8 VIPE_EXT_JIT=1 MAX_JOBS=8 uv run python -m unittest tests.test_graph_buffer_fused_ba -v
  • CUDA_HOME=/usr/local/cuda-12.8 VIPE_EXT_JIT=1 MAX_JOBS=8 uv run python -m unittest tests.test_slam_ext_ba_extended -v
  • uv run pytest tests/test_config.py -q
  • uv run ruff format --check vipe/config/slam.py vipe/slam/components/buffer.py tests/test_config.py tests/test_graph_buffer_fused_ba.py tests/test_slam_ext_ba_extended.py
  • uv run ruff check vipe/config/slam.py vipe/slam/components/buffer.py tests/test_config.py tests/test_graph_buffer_fused_ba.py tests/test_slam_ext_ba_extended.py
  • uv run mypy vipe/config/slam.py vipe/slam/components/buffer.py tests/test_config.py tests/test_graph_buffer_fused_ba.py

@heiwang1997
heiwang1997 merged commit 536fa8f into main May 22, 2026
1 check passed
@heiwang1997
heiwang1997 deleted the code/fused-ba branch May 22, 2026 00:16
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