Skip to content

perf: optimize sample_floyd by unsafe APIs #1622

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

perf: optimize sample_floyd by unsafe APIs #1622

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
07d4e92
Boost performance for sample_floyd
Unparalleled-Calvin Mar 30, 2025
5e03158
Merge branch 'master' into master
Unparalleled-Calvin Apr 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ You may also find the [Upgrade Guide](https://rust-random.github.io/book/update.
- Fix feature `simd_support` for recent nightly rust (#1586)
- Add `Alphabetic` distribution. (#1587)
- Re-export `rand_core` (#1602)
- Boost performance of `sample_floyd` (#1622)
- Allow `fn rand::seq::index::sample_weighted` and `fn IndexedRandom::choose_multiple_weighted` to return fewer than `amount` results (#1623), reverting an undocumented change (#1382) to the previous release.

## [0.9.0] - 2025-01-27
Expand Down
18 changes: 13 additions & 5 deletions src/seq/index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

//! Low-level API for sampling indices
use alloc::vec::{self, Vec};
use core::ptr;
use core::slice;
use core::{hash::Hash, ops::AddAssign};
// BTreeMap is not as fast in tests, but better than nothing.
Expand Down Expand Up @@ -447,12 +448,19 @@ where
// the last entry. This bijection proves the algorithm fair.
debug_assert!(amount <= length);
let mut indices = Vec::with_capacity(amount as usize);
for j in length - amount..length {
let t = rng.random_range(..=j);
if let Some(pos) = indices.iter().position(|&x| x == t) {
indices[pos] = j;
let mut len = 0;
let ptr = indices.as_mut_ptr();
// safety: the index is bounded by the length of indices
unsafe {
for j in length - amount..length {
let t = rng.random_range(..=j);
if let Some(pos) = indices.iter().position(|&x| x == t) {
*indices.get_unchecked_mut(pos) = j;
}
Comment on lines +457 to +459
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we need to use an unsafe function here at all. Try re-writing the iterator with a simple for loop:

for pos in 0..indices.len() {
    if indices[pos] == t {
        indices[pos] = j;
        break;
    }
}

Sure, that uses index operations but the compiler should be able to remove the range checks. (I'm a little surprised if it can't here.)

Copy link
Author

@Unparalleled-Calvin Unparalleled-Calvin Apr 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Emm, I try to rewrite this part with for loop. The benchmark result shows that it has similar performance to the original one.. It seems that the part optimized by unsafe APIs can not be achieved by the compiler easily.

ptr::write(ptr.add(len), t);
len += 1;
indices.set_len(len);
Comment on lines +460 to +462
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this actually any faster than push? I suppose it may be (eliminating the capacity check).

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes for sure :)

}
indices.push(t);
}
IndexVec::from(indices)
}
Expand Down