Skip to content

Rust acceleration via PyO3/maturin#5

Merged
craftsangjae merged 8 commits into
mainfrom
feat/rust-acceleration
Feb 16, 2026
Merged

Rust acceleration via PyO3/maturin#5
craftsangjae merged 8 commits into
mainfrom
feat/rust-acceleration

Conversation

@craftsangjae
Copy link
Copy Markdown
Owner

@craftsangjae craftsangjae commented Feb 16, 2026

Summary

  • Accelerate core operations (Gray code partitioning, scatter-add, empty partition filling) in Rust via PyO3/maturin
  • Switch build system from hatchling to maturin with cross-platform wheel builds in CI
  • Graceful fallback to pure Python when the Rust extension is unavailable

Benchmark (ARM64, Python 3.9, NumPy 2.0)

Single document encoding (128 vectors, dim=128):

Config Rust Python Speedup
R=20, k=5, d_proj=8 0.69 ms 11.51 ms 16.7x
R=20, k=5, d_proj=16 0.74 ms 12.00 ms 16.2x

Batch document encoding (100 docs, ~128 vectors each):

Config Rust Python Speedup
R=20, k=5, d_proj=8 17.89 ms 32.07 ms 1.8x
R=20, k=5, d_proj=16 23.42 ms 57.66 ms 2.5x

Changes

  • New: `Cargo.toml`, `src/{lib,gray_code,partition,scatter,fill_empty}.rs` — Rust kernel implementations
  • New: `muvera/_rust_kernels.pyi` — type stubs for the Rust module
  • New: `tests/test_rust_equivalence.py` — 13 numerical equivalence tests (Rust vs Python)
  • Modified: `pyproject.toml` — switch to maturin build backend
  • Modified: `muvera/{init,helper,muvera}.py` — call Rust kernels with graceful fallback
  • Modified: `.github/workflows/{test,publish}.yml` — add Rust toolchain + maturin-action for wheel builds
  • Modified: `README.md` — add Rust acceleration benchmark results
  • Modified: `examples/basic_usage.py` — use `list[np.ndarray]` input format

Test plan

  • All 70 existing tests pass unchanged
  • 13 new Rust/Python numerical equivalence tests pass
  • `ruff check .` and `ruff format --check .` pass
  • `examples/basic_usage.py` runs successfully
  • Benchmark measurements completed

Signed-off-by: craftsangjae <craftsangjae@gmail.com>
Signed-off-by: craftsangjae <craftsangjae@gmail.com>
…tion in the MuVERA Python implementation.

Signed-off-by: craftsangjae <craftsangjae@gmail.com>
maturin develop requires a virtualenv which is not available
in GitHub Actions. pip install invokes maturin as the build
backend automatically.
… binary

Signed-off-by: craftsangjae <craftsangjae@gmail.com>
@craftsangjae craftsangjae self-assigned this Feb 16, 2026
@craftsangjae craftsangjae merged commit d90422f into main Feb 16, 2026
6 checks passed
@craftsangjae craftsangjae deleted the feat/rust-acceleration branch February 16, 2026 05:30
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