implements weighted shuffle using N-ary tree

[behzadnouri]

This is port of firedancer's implementation of weighted shuffle:
https://github.com/firedancer-io/firedancer/blob/3401bfc26/src/ballet/wsample/fd_wsample.c

Problem

#185 implemented weighted shuffle using binary tree. Though asymptotically a
binary tree has better performance, compared to a Fenwick tree, it is
less cache local resulting in smaller improvements and in particular
slower WeightedShuffle::new.

Summary of Changes

In order to improve cache locality and reduce the overheads of
traversing the tree, this commit instead uses a generalized N-ary tree
with fanout of 16, showing significant improvements in both
WeightedShuffle::new and WeightedShuffle::shuffle.

With 4000 weights:

N-ary tree (fanout 16):

test bench_weighted_shuffle_new     ... bench:      36,244 ns/iter (+/- 243)
test bench_weighted_shuffle_shuffle ... bench:     149,082 ns/iter (+/- 1,474)

Binary tree:

test bench_weighted_shuffle_new     ... bench:      58,514 ns/iter (+/- 229)
test bench_weighted_shuffle_shuffle ... bench:     269,961 ns/iter (+/- 16,446)

Fenwick tree:

test bench_weighted_shuffle_new     ... bench:      39,413 ns/iter (+/- 179)
test bench_weighted_shuffle_shuffle ... bench:     364,771 ns/iter (+/- 2,078)

The improvements become even more significant as there are more items to
shuffle. With 20_000 weights:

N-ary tree (fanout 16):

test bench_weighted_shuffle_new     ... bench:     200,659 ns/iter (+/- 4,395)
test bench_weighted_shuffle_shuffle ... bench:     941,928 ns/iter (+/- 26,492)

Binary tree:

test bench_weighted_shuffle_new     ... bench:     881,114 ns/iter (+/- 12,343)
test bench_weighted_shuffle_shuffle ... bench:   1,822,257 ns/iter (+/- 12,772)

Fenwick tree:

test bench_weighted_shuffle_new     ... bench:     276,936 ns/iter (+/- 14,692)
test bench_weighted_shuffle_shuffle ... bench:   2,644,713 ns/iter (+/- 49,252)

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.9%. Comparing base (a916edb) to head (55d94d1).

Additional details and impacted files

@@            Coverage Diff            @@
##           master     #259     +/-   ##
=========================================
- Coverage    81.9%    81.9%   -0.1%     
=========================================
  Files         840      840             
  Lines      228105   228102      -3     
=========================================
- Hits       186837   186832      -5     
- Misses      41268    41270      +2     

[gregcusack]

those performance numbers look great! why did you settle on FANOUT of 16? just performed the best when you tested or?? Looks like firedancer is using 9. either way lgtm!

[behzadnouri]

why did you settle on FANOUT of 16? just performed the best when you tested or??

yeah, ran the benchmarks with different values of BIT_SHIFT and 4 seemed fastest.

[mergify]

Backports to the beta branch are to be avoided unless absolutely necessary for fixing bugs, security issues, and perf regressions. Changes intended for backport should be structured such that a minimum effective diff can be committed separately from any refactoring, plumbing, cleanup, etc that are not strictly necessary to achieve the goal. Any of the latter should go only into master and ride the normal stabilization schedule. Exceptions include CI/metrics changes, CLI improvements and documentation updates on a case by case basis.