Rust: Use rayon and better performance #203
Conversation
|
The issue with Rayon or any dependencies is that it brings a lot of new dependencies that need to be audited or at least vetted. This is very undesirable in a cryptographic product securing billions of assets. It is especially critical when supply chain attacks are increasing in appeal. Furthermore, a simple threadpool written with cryptographic/security application in mind and fully control by the dev is significantly easier to maintain and update if there is a security issue. |
I generally agree with you on this one. However, as I have shown in the description already, there were some questionable decisions from soundness point of view around thread pool as it was used and performance was lower than it could have been. You 100% can rewrite original code and make it faster without rayon, achieving rayon's features in a safe misuse-resistant way is non-trivial and I'd rather use something that it maintained by Rust experts for this purpose (check top contributors to rayon project, they are literally working on Rust itself) than seeing every library invent their own a bit flawed version of the same features that doesn't integrated as well with the rest of Rust ecosystem. I would estimate risk of relying on |
| @@ -25,8 +25,7 @@ macro_rules! pippenger_test_mod { | |||
| let mut scalars = Box::new([0u8; nbytes * npoints]); | |||
| ChaCha20Rng::from_seed([0u8; 32]).fill_bytes(scalars.as_mut()); | |||
|
|
|||
| let mut points: Vec<$point> = Vec::with_capacity(npoints); | |||
| unsafe { points.set_len(points.capacity()) }; | |||
| let mut points = vec![<$point>::default(); npoints]; | |||
There was a problem hiding this comment.
ditto: avoid useless zero init
| @@ -61,8 +60,7 @@ macro_rules! pippenger_test_mod { | |||
| let mut scalars = Box::new([0u8; nbytes * npoints]); | |||
| ChaCha20Rng::from_seed([0u8; 32]).fill_bytes(scalars.as_mut()); | |||
|
|
|||
| let mut points: Vec<$point> = Vec::with_capacity(npoints); | |||
| unsafe { points.set_len(points.capacity()) }; | |||
| let mut points = vec![<$point>::default(); npoints]; | |||
There was a problem hiding this comment.
ditto: avoid useless zero init
There was a problem hiding this comment.
Thanks for review @mratsim! If the rest of the stuff is fine, I can try to remove zero initialization and see if it makes a performance difference, if yes then I'll look to refactor things in a way that makes code easier to audit.
Having documentation at least in C code would have been great, ideally in bindings as well. Right now without C expanding macros in my head I have no idea that blst_p1s_to_affine initializes memory internally and is guaranteed to never ever read it prior to that, which to me is horrifying from security point of view, especially when there is an easy way to avoid it.
… newer version of Rust
…dependency and unnecessary/unused `no-threads` feature
There was a problem hiding this comment.
I refactored blst_fp12::miller_loop_n into something that is close to original implementation.
For zero initialized vectors I wrote this simple bench:
fn bench_alloc(c: &mut Criterion) {
let mut group = c.benchmark_group("allocation");
group.bench_function("alloc/set_size", |b| {
b.iter(|| {
let mut points =
Vec::<u8>::with_capacity(criterion::black_box(1_000_000));
unsafe { points.set_len(points.capacity()) };
criterion::black_box(points);
});
});
group.bench_function("alloc/zeroed", |b| {
b.iter(|| {
let points = vec![0u8; criterion::black_box(1_000_000)];
criterion::black_box(points);
});
});
group.finish();
}The results were like this:
allocation/alloc/set_size
time: [12.233 ns 12.276 ns 12.321 ns]
allocation/alloc/zeroed time: [11.259 µs 11.263 µs 11.267 µs]
It seemed significant, so I decided to partially revert set_len, but I moved some of them after the function that initializes memory to make it clearer to the reader what is actually happening there. Only left zero initialization in 2 places in tests where performance doesn't matter as much.
And I applied relaxed ordering in places where suggested.
@mratsim, appreciate careful review, hopefully you'll have time to check the diff.
| @@ -25,8 +25,7 @@ macro_rules! pippenger_test_mod { | |||
| let mut scalars = Box::new([0u8; nbytes * npoints]); | |||
| ChaCha20Rng::from_seed([0u8; 32]).fill_bytes(scalars.as_mut()); | |||
|
|
|||
| let mut points: Vec<$point> = Vec::with_capacity(npoints); | |||
| unsafe { points.set_len(points.capacity()) }; | |||
| let mut points = vec![<$point>::default(); npoints]; | |||
| @@ -61,8 +60,7 @@ macro_rules! pippenger_test_mod { | |||
| let mut scalars = Box::new([0u8; nbytes * npoints]); | |||
| ChaCha20Rng::from_seed([0u8; 32]).fill_bytes(scalars.as_mut()); | |||
|
|
|||
| let mut points: Vec<$point> = Vec::with_capacity(npoints); | |||
| unsafe { points.set_len(points.capacity()) }; | |||
| let mut points = vec![<$point>::default(); npoints]; | |||
|
Just confirmed that performance improvements here translate to performance improvements downstream in our usage usage of https://github.com/sifraitech/rust-kzg (uses blst) in https://github.com/subspace/subspace, both of which also take advantage of Rayon as well. |
I don't know whether commitment, proof generation or verification is the bottleneck for your use-case but @sauliusgrigaitis might be able to point you on how to best use rust-kzg, i.e. either through native BLST or reimplemented+parallelized BLST (i.e. using their own multithreading strategy). Also they've been working on a new set of benchmarks. I have my own in mratsim/constantine#304 (comment), but I don't cover parallel BLST. |
|
@dot-asm is there something I can do to interest you reviewing this? |
This PR introduces a bunch of changes that target at better performance and higher code quality.
Performance
Before:
After:
The reason is most likely due to higher parallelism and not using channels in those operations.
Rayon
Rayon doesn't always result in better performance than manual thread management (see rayon-rs/rayon#1124), but changes in this PR do result in higher performance.
There are many great benefits to rayon though in the way it is used:
safe API, previous thread pool usage was either unsound or sound by accident (comments were claiming things like:
Which while worked by accident, was a lie (nothing was actually joined in most cases, but worked in practice due to use of channels to send work back to the parent thread), fundamentally unsound API (parent thread panic would result in worker threads accessing undefined memory location) and generally a major foot gun for anyone who would try to change the code in any way.
support for customizing number of threads using environment variables
support for customizing number of threads and their pinning to CPU cores depending on thread pool from which operations are called (by default global thread pool is used and default behavior is equivalent to old code)
avoid conflicting/overlapping with other threads if used from within parallel code that already uses rayon
no need to special-case platforms like WASM, if specific WASM environment supports threading,
blstwill now enjoy it automaticallyVarious cleanups
I changed all but one
channelusage with a cleaner code usingMutex,Arcusage was removed as unnecessary and some other bookkeeping was cleaned up after understanding what code does and rewriting intention in more idiomatic Rust code (seefn fromfor example). The only place wherechannelstill remains isfn mult, but it is quite difficult code and I didn't have the will to decipher what it tries to do to get rid of channel, but hopefully examples in this PR will prompt someone to clean it up at some later point.Some
unsafeblocks withvec.set_len()are moved after memory is actually initialized and some questionable in terms of soundness places were refactored such that they are quite obviously sound now. Some remaining places can also be fixed, but require pointer arithmetic that might require newer versions of Rust, so I left TODOs there.I also took liberty to refactor things like
&q[0]toq.as_ptr()since while both result in the same pointer, the former is deceptive because it suggests to the reader that only the first element is used, while in practice the wholeqis used by the function that is being called later (documentation is almost non-existent, so it wasn't obvious at first and I hope this improves readability).no-threads
no-threadswas a hack and non-idiomatic because it violates "additivity" assumptions of Rust features. Thankfully it is no longer necessary and can be safely removed, but since this is a breaking change (the only one in this PR to the best of my knowledge), I bumped crate version from0.3.11to0.4.0.How to review this PR
This PR is split into self-contained logically structured commits and it is recommended to review commits one by one instead of the complete diff, this way it is much easier to understand what was changed and why.
Let me know if there are any questions, happy to explain each line change in detail.