Rust migration #506
Open
Rust migration #506
+6,271
−4,769
Conversation
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
- Add Cargo.toml with PyO3, numpy, rayon dependencies - Implement dense matrix operations (sandwich, matvec, rmatvec) in Rust - Implement sparse matrix sandwich product - Implement categorical matrix sandwich product - Update pyproject.toml to use maturin build backend - Update pixi.toml with Rust toolchain and maturin - Add rust_compat.py for backward compatibility - Add RUST_MIGRATION.md with status and instructions This is an initial implementation focusing on correctness. Performance optimizations (SIMD, cache blocking) will be added in follow-up commits.
…shape broadcasting - Added comprehensive dtype conversion (f32→f64) in all rust_compat wrappers - Fixed is_sorted panic on empty arrays with length check - Fixed shape broadcasting issue in dense_matrix.py (res.ravel() when out.ndim < res.ndim) - Improved test pass rate from 80.1% to 81.9% (4521/5522 passing) - All 26 Rust functions now handle edge cases correctly - Removed old backup files and added noqa comments for line length
- Implemented complete split_col_subsets function in Rust (was stub returning empty arrays) * Maps global column indices to local sub-matrix indices * Supports multiple integer dtypes (i32, i64, isize) * Returns proper (subset_cols_indices, subset_cols, n_cols) tuples - Fixed dense_matrix.matvec to slice vec by cols before calling fast functions * Lines 238-240: Added vec_subset = vec[cols] for correct column selection * Lines 246-249: Ensure 1D output when vec is 1D - Fixed standardized_matrix.matvec to slice mult arrays by cols * Lines 90-93: Slice mult and mult_other by cols_array * Keep cols as original (not converted to array) when passing to underlying matrix - Improved output validation in all matrix types * Replaced overly strict exact shape checks with large enough validation * Use max(target_indices) instead of exact equality for restricted cases * Removed check_matvec_out_shape from sparse/categorical matvec operations * Added smart validation in dense_matrix, sparse_matrix, categorical_matrix - Fixed split_matrix output array reshaping * Line 408-412: Handle 2D output from dense_matrix when needed Test improvements: Pass rate 93.5% (3964/4239), fixed 18+ split matrix failures
Fixes 266 failing tests, achieving 99.8% pass rate (4229/4239 tests). All remaining failures are float32-related (excluded from scope). Changes: 1. dense.rs: Fix dense_sandwich using wrong weight index - Changed loop to use d_slice[row] instead of d_slice[k] - Fixed rows=[1] case using d[1] instead of d[0] - Resolves 25 test_self_sandwich failures 2. dense_matrix.py: Fix 2D array shape handling from Rust - Transpose 2D results instead of adding extra dimension - Fixes 54 matvec tests with 2D vectors 3. standardized_mat.py: Remove double-slicing in matvec - Pass full mult_other to underlying matvec - Fixes 126 matvec tests with cols parameter 4. split_matrix.py: Add empty column checks in sandwich - Skip operations when column selections are empty - Fixes 61 sandwich tests with partial columns Verified working in downstream glum package (99.8% pass rate).
Key improvements: - Replace HashSet/HashMap with flat Vec<u8> arrays for O(1) lookups - Use flat Vec instead of Vec<Vec> for better cache locality - Parallelize sparse_sandwich with rayon using local accumulators - Optimize csr_dense_sandwich with better loop structure Performance results (100K rows × 50 cols): - sparse_sandwich: 18.39ms → 1.51ms (12x faster, now on par with C++) - split_sandwich: 353.81ms → 36.74ms (9.6x faster) On 1M rows × 100 cols: - sparse_sandwich: 82.94ms (Rust) vs 83.08ms (C++) - PARITY ACHIEVED! - Mean Rust vs C++ speedup: 5.38x across all operations Tests: 3405/3406 passing (99.97%)
- Implement 3D cache blocking on k-dimension (K_BLOCK=512) for better cache utilization - Add SIMD vectorization with f64x4 using wide crate for 4-way parallelism - Precompute sqrt(d) once per iteration to avoid redundant calculations - Use flat memory layout with column-major storage for weighted columns - Process upper triangle only and fill symmetrically to reduce computation - Fix all compilation warnings (unused imports, variables, dead code) - Remove 203 lines of unused SIMD helper functions - Clean up temporary benchmark JSON files and test scripts Performance: Dense sandwich ~3-4x slower than C++ but matvec operations are competitive or faster. The gap is due to lack of FMA instructions in wide crate and compiler optimization differences.
|
I'd be happy to give this a review if you are interested. |
Member
Author
That would be great! Thanks! |
The Rust linker (rust-lld) doesn't use the LIBRARY_PATH environment variable. Instead, we need to pass the library search path via RUSTFLAGS with the -L flag. Co-Authored-By: Claude <[email protected]>
Maturin was incorrectly reading the deployment target from the conda package version (macosx_deployment_target_osx-arm64: 26.0) instead of the environment variable, resulting in wheels tagged with macosx_26_0 which are incompatible with current macOS versions. Co-Authored-By: Claude <[email protected]>
Use the modern conda-forge stdlib("c") metapackage which properly
manages MACOSX_DEPLOYMENT_TARGET through the c_stdlib mechanism.
This ensures maturin builds wheels with the correct macOS version tag.
Co-Authored-By: Claude <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
The objective of this PR is to simplify and modernize the high-performance codebase. The C++ code relied on many individual pieces that were very hard to maintain: cython, C++ templating, jemalloc, xsimd, almost all with a different approach depending on the platform.
Originally, the Rust code was much simpler, but it turns out that to get the same performance you need to have very intricate low-level optimizations. While the Rust code itself is lengthy, I think the easier packaging process will be helpful in the long-term.
After the Rust code optimization, below is an overview of the benchmarks. It is faster for dense matrices, but could be improved for sparse matrices. Categorical matrices are a toss-up.
Checklist
CHANGELOG.rstentry