FastJet but faster? Part II
some fun Rust optimization lessons
tldr
We never set out to make fastjet-rs faster than FastJet.
But after some elbow grease, SIMD, and some not-so-micro optimizations we’ve managed to make our Rust simple jet finding algorithms a few microseconds faster than the original C++ implementations.
If you haven’t already gone over the basics of this project in part 1, all you need to know is that fastjet is a particle physics package that performs nearest neighbor calculations to merge subatomic particles into jets. We are trying to rewrite the original C++ as Rust for the language’s memory safety advantages.
The benchmarking results from a comparison of the simple O(n²) clustering algorithm found that our Rust implementation is roughly equivalent in performance to a Julia implementation, JetReconstruction.jl, but ultimately slower than the original C++ version by ~130 microseconds… but that’s not the case anymore!
Jump to the end if you’re just interested in the results of our benchmarking.
We’ll go over three major categories of code changes since our commit as of Part 1.
tiling
Thanks JetReconstruction.jl team for this wonderful graphic!
Fastjet and JetReconstruction.jl both implement a tiled clustering algorithm in addition to the simple n² clustering algorithm so we thought we should too. The key to this algorithm lies in the fact that clustering particles on an eta, phi coordinate system doesn’t change too much if we localize the clustering. In other words, instead of scanning and calculating distances between one jet and all other jets, we can define a tiled representation of our particles so that we only use the closest particles during distance calculations. Depending on how small we make our tiles, we essentially do the minimum amount of distance calculations for each jet at the expense of extra bookkeeping as we now need to store which jets are in which tiles and actively update it.
In practice, this means implementing a doubly linked list.
Let’s introduce some new structs Tile and TiledJet to represent all of our jet information relationally.
/// The neighbourhood of a tile is set up as follows:
/// LRR
/// LXR
/// LLR
///
/// tiles contains XLLLLRRRR with pointers
/// | \ RH_tiles
/// \ surrounding_tiles
Credit: fastjet authorsChoosing a doubly linked list gives us instant access to adjacent jets in both directions, making it easy to rescan neighbors around a merge. Flattening each tile makes traversal linear, all we have to do is a little bit of math to map the grid index of each neighboring jet in a tile to its flattened index. Looks pretty easy to implement right?
why can doubly linked lists be unsafe in Rust?
Linked lists are a bit of a pain to implement in Rust.
struct Node {
val: u32,
prev: Option<Box<Node>>,
next: Option<Box<Node>>
}Since each node needs to hold references to both the prev node and the next node, any interior node will always have multiple simultaneous owners. The moment you try to mutate a node in the middle of the list the Rust borrowchecker will throw a fit since you’ll need to have write permissions to both the node itself, the previous node, and the next node.
Since each node needs to hold references to both the prev node and the next node, any interior node will always have multiple simultaneous owners. The moment you try to mutate a node in the middle of the list the Rust borrowchecker will throw a fit since you’ll need to have write permissions to both the node itself, the previous node, and the next node.
why RefCell is ‘essential’
Recall RefCell<T> from part 1, paired with Rc<T> they are the Rust sanctioned safe method of allowing shared mutability. If you wrap nodes with Rc<RefCell<Node>> ownership checks will now be performed at runtime instead of compile time. This is significant since the borrow checker, which runs at compile time is quite conservative so it ends up rejecting any borrowing patterns that can’t be proven safe statically. At runtime all of this changes since we can actually check if any simultaneous borrows conflict. Thus our doubly linked list node implementation looks more like this:
struct Node {
val: u32,
prev: Option<Rc<RefCell<Node>>>,
next: Option<Rc<RefCell<Node>>>
}This will compile successfully, but the problem comes when actually accessing any of our jets inside the linked list. Each time you want to read a field within an object of type Node you will need to call .borrow() or .borrow_mut(). Each call to either of these functions is a runtime borrow check, adding a significant overhead anytime we work with jets in our algorithm.
ok but what if we don’t want to always have runtime borrow checking?
A better solution is actually pretty elegant here if you are willing to trade in storage overhead for better runtime performance. The astute reader might have already noticed that we don’t actually need to store references to each jet object, we can just keep a well updated table of jets and use usize indices instead.
Our main clustering logic is implemented inside ClusterSequence so we can give it sole ownership over this table. This is actually a common design pattern more commonly known as arena-style bookkeeping. Instead of using pointers or references, we can store all of our objects in contiguous blocks of memory that we access via indices. However, since our clustering algorithm requires we merge quite a few jets the upkeep of marking jet indices stale was nontrivial in our implementation. Many bugs emerged due to an erroneous index to an already removed or merged jet. All of this still ended up being worth it in our case since our runtime performance did improve after removing several repeated runtime borrow checks.
So linked lists really aren’t all that easy to pull off in Rust without giving some thought to ownership models in your data structures. The reason why implementing linked lists might feel easier in other languages is simply because the memory management is abstracted away behind complex garbage collection cleanup that has its own hidden costs. From a pedagogical perspective, tiling our clustering algorithm turned out to be a great way to learn Rust.
For a more thorough explanation of all things linked lists in Rust we highly recommend checking out Learning Rust With Entirely Too Many Linked Lists :)
the fallacy of SIMD
Implementing a SIMD version of the simple n² clustering algorithm was the second design change we made to increase our performance.
SIMD stands for (Single Instruction Multiple Data) where the CPU groups together multiple data elements as vectors (defined as LANES) and performs the same instruction to all this data simultaneously. In Fastjet, there are multiple scenarios where our data is highly parallel, finding the distance metric for all jets at once, initializing them, and updating the nearest neighbors each time we remove/add a new jet. Some other parts of the code are inherently sequential, such as finding which jet to remove and performing our updating logic on it since it relies on finding the minimum distance jet and then acting upon it. In these cases the latter is more ripe for SIMD than the former. In these cases SIMD can optimistically produce performance gains of 3–4x, however, the compiler is a fickle thing and it is possible that SIMD can cause performance penalties because nothing beats the compiler.
Normally, SIMD requires fine tuning to specific CPU architectures and calling archaically named functions after scouring through CPU documentation. Thankfully for our purposes we can test drive std::simd, a feature of Rust nightly which has long awaited a stable release, that aims to provide a “universal” SIMD implementation. The initial setup was pretty straightforward other than actually needing to learn SIMD!
When it came to retrofitting our code for SIMD we faced a number of difficulties.
SIMD memory requirements
Ref: https://www.researchgate.net/figure/A-comparison-of-the-Array-of-Structs-AoS-and-Struct-of-Arrays-SoA-memory-layouts_fig2_259132352
First off, SIMD requires all the data to be contiguously located in memory, meaning that we had to vectorize our current paradigm where each Jet struct contains elements containing its pt, eta, phi, NN_index, etc. This means that every time we call a SIMD function we need to loop over all vectors and initialize arrays holding just phi, eta, and so on. Initializing this data each SIMD call would increase the overhead so much that the performance gain is unjustified, thus we must shift to running our entire algorithm using structures of arrays instead of arrays of structs.
For this reason, we were drawn to Rusts’ arena (recall the importance of arena-style bookkeeping from earlier!) since it allows for fast memory allocation given that we know the amount of Jets that we account for at a time will never be larger than the initial length of input particles. This required a large overhaul of our code, but it was all going to be worth it right?
SIMD and primitives
Another issue we ran into was the fact that SIMD arrays only take primitives but our NN_index was encoded as an Option (to highlight jets not having NN) could not be imputed. Then thanks to (JetReconstruction.jl) we realized that we can actually modify NN_index to be set to a sentinel value of itself!
When creating new jets (whether merged or initialized) we already know its own index to access it, and most functions where we unwrap can instead just check its index. This mainly improves performance (removing unwrapping and wrapping) but also has the secondary effect of reducing data size, f64 is 8 byte while an Option is double that (just to store the None value).
An interesting example to highlight was the following function since the min(a, a) = a:
fn _bj_dij<J: ProxyJet>(jet: &J, jets: &[J]) -> f64 {
jet.nn_dist()
* jet.kt2().min(
jet.nn_jet_index()
.map(|index| jets[index].kt2())
.unwrap_or(f64::MAX),
)
}
fn _bj_dij<J: ProxyJet>(jet: &J, jets: &[J]) -> f64 {
let kt2_a = jet.kt2();
let kt2_b = jets[jet.nn_jet_index()].kt2();
// Bypasses the f64::min NaN checks, yielding pure hardware speed
let min_kt2 = if kt2_a < kt2_b { kt2_a } else { kt2_b };
jet.nn_dist() * min_kt2
}Now with all the setup out of the way we can examine a smaller SIMD function to highlight its technicalities. Since we showed the _bj_dij function earlier it seems right to highlight it parallelized.
#[cfg(feature = “simd”)]
fn _dij_simd<const LANES: usize>(
kt2_arr: &[f64],
nn_idx: &[u64],
nn_dist: &[f64],
di_j: &mut [f64],
) {
let n = kt2_arr.len();
let mut j = 0;
// SIMD loop
while j + LANES <= n {
// note: [j..j + LANES] does not include j + LANES
let kt2_i = Simd::<f64, LANES>::from_slice(&kt2_arr[j..j + LANES]);
let dist_i = Simd::<f64, LANES>::from_slice(&nn_dist[j..j + LANES]);
let indices = Simd::<u64, LANES>::from_slice(&nn_idx[j..j + LANES]);
// must store indices as u64 because compiler errors with other funcs
let indices_usize = indices.cast::<usize>();
let kt2_b = Simd::<f64, LANES>::gather_or(kt2_arr, indices_usize, kt2_i);
let kt2_min = kt2_i.simd_min(kt2_b);
let res = kt2_min * dist_i;
res.copy_to_slice(&mut di_j[j..j + LANES]);
j += LANES;
}
// scalar handling for remainder of elems aka tail
for i in j..n {
let idx = nn_idx[i] as usize;
di_j[i] = nn_dist[i] * kt2_arr[i].min(kt2_arr[idx]);
}
}Starting from the top, we gate SIMD use only when building with feature “simd” which prevents compiler errors from not being on nightly Rust as the code is basically “ignored” by the preprocessor.
Then when calling these SIMD functions we must pass a size parameter labeled LANES which as mentioned previously is how many elements we wish to pass to a SIMD instruction at a time. It is up to the user to check whether they handle the remainder of elements with the slower sequential instruction which we do in the final loop at the end. We can construct SIMD objects easily using from_slice and just passing a reference to our elements. Here the parallel action we must do is gather_or which retrieves all elements given indices. Funnily enough this is one of the slower operations since the bottleneck for these operations are not the CPU but instead the memory register which must fetch the vectors which are not contiguous and not guaranteed to be in the same cache line as the jets we have already taken in. However, we can see how we can use many other primary operators like inequalities, addition, multiplication in the same way and chain these operations together.
After overhauling all of our code we found that before optimizing our simple_n2_cluster algorithm that we achieved a 45 percent performance improvement, and we were officially slightly faster than our C++ “competitor” 🥳.
micro optimizations that are more than micro in practice
However, it turns out that just a couple of changes in our previous naive implementation that help the compiler do it’s job also did wonders for performance.
writing clean code really does count!
Part of the reason that Rust touts ‘C’ performance with large safety guarantees is through its aggressive use of optimizations to construct Zero-Cost Abstractions. This means more verbose Rust code to check indices in comparison to unsafe pointers, Vectors instead of arrays compile down to almost the same assembly instruction as C’s unsafe implementation.
This is important to note as more convoluted code has the effect of preventing the compiler from making such optimizations since user intent is unknown. This is why inlining most functions can result in faster performance since the compiler knows in what situations called functions are being used by the callee and can optimize accordingly. Take the following for loop:
for i in 0..tail {
// some logic based on the value of nn_idx that only matters if jet_b != i
// if we find jet_b_idx is equal to i quit
if jet_b_idx == i {
if jets[jet_b_idx].nn_jet_index() == tail {
jets[jet_b_idx].set_nn_jet(min_idx);
}
continue;
}
// update min dist for jet i and jet_b
}Given this function the compiler does not know that the if loop is only checking for one value, and one that is predetermined before the for loop even starts, since it is in the middle of the function. If we just place it on top, the compiler will split the function into the following:
for i in 0..jet_idx_b {
// do logic
}
// perform jet_b_idx == i case
for i in jet_idx_b+1..tail {
// do logic
}A small change like moving the if loop to the beginning of the for loop had around a 6 percent improvement since this was a part of the code that was called very frequently with depth O(n²).
performance impact of muts
The next example highlighted to me why muts can actually decrease performance if not used sparingly especially in conjunction with the next example. When a jet is merged, we must set the NN for the old jets which previously pointed to its constituents, since we don’t have to update the other jets to point to it (since we removed them) as a NN we only need to update that jet. Here is the code to highlight this.
#[inline]
fn bj_set_nn_nocross<J: ProxyJet>(
&self,
curr_idx: usize,
head_idx: usize,
tail_idx: usize,
jets: &[J],
) {
let mut nn_dist = self.r2;
let mut nn: usize = curr_idx;
for jet_b_idx in head_idx..curr_idx {
let dist = J::_bj_dist(curr_jet, &jets_ref[jet_b_idx]);
if dist < nn_dist {
nn_dist = dist;
nn = jet_b_idx;
}
}
for jet_b_idx in curr_idx + 1..tail_idx {
let dist = J::_bj_dist(curr_jet, &jets_ref[jet_b_idx]);
if dist < nn_dist {
nn_dist = dist;
nn = jet_b_idx;
}
}
jets[curr_idx].set_nn_jet(nn);
jets[curr_idx].set_nn_dist(nn_dist);
}The important things to watch are the function signature, and the last two lines. Since we need to modify jets[curr_idx] we pass the entire jets as mutable, we also can’t just pass jets[curr_idx] as mutable because then we would break Rust’s borrowing rules. Thus the only solution is to tell the callee function which element changes, and then modify that element. Here is how that following code looks like:
let (new_nn, new_dist) = self.bj_set_nn_nocross(i, 0, tail, jets);
let jet_i = &mut jets[i];
jet_i.set_nn_jet(new_nn);
jet_i.set_nn_dist(new_dist);
fn bj_set_nn_nocross<J: ProxyJet>(
&self,
curr_idx: usize,
head_idx: usize,
tail_idx: usize,
jets: &[J],
) -> (usize, f64) {
// logic
(nn, nn_dist)
}This allows us to avoid having to pass the entire jets collection as mutable preventing the need for the compiler to reload the entire jets collection after passing into the function. From the compiler’s eyes the entire vector could change. With this declaration we know that only jets[i] changes. This change resulted in a 6 percent performance improvement.
eliding compiler bound checks
This last example is arguably the most important. Each time a vector element is indexed, Rust checks whether the element is out of bounds (OOB) which in C leads to a host of UB. Given that our entire code relies on NN_index and indexing our jets almost every time we require them, this becomes a large part of our performance penalty.
Iterator efficiency in Rust can also be explained by the same phenomenon, since iterators already check elements exist, the compiler does not need to perform auxiliary bounds checking!
For us, we can’t use iterators here because we need to pass the entire vector slice to internal functions alongside a jet from that slice. This doesn’t mean we have no options left though. Instead we can tell the compiler to stop bounds checking indexes by proving that we will never use an index that is OOB via assertions.
In the code highlighted in the first optimization example, there is a lot of indexing. Almost everything is being indexed since we also have muts sprinkled in the for loop. Thus before this entire loop we index according to i and jet_b_idx thus we add the following code before the loop.
let jets = &mut bj_jets[..tail];
let dijs = &mut di_j[..tail];
assert!(jet_b_idx < jets.len());Since i loops up to tail, and jet_b is less than the entire length of the jet, OOB cannot occur and the compiler removes all bounds checking for our code in this for loop and in our NN search due to the changes mentioned in the second example. This led to an astonishing 20% percent from just 3 lines of code!
Other not so micro optimizations that accounted for another 20% of performance improvements:
- Knowing our elements were never
NaNorINFwhich require another check forf64when sorting or min - Removing lazy init since just calculating these numbers are cheaper than all of the dereferencing and heap retrieval
- Using Fused
add_mulin math calculations
Lastly, we found that testing different performance improvements with Rust was fairly easy as we didn’t have to worry about UB caused by our code changes that would normally be present in C++ pointer handling.
so how much faster are we actually?
We used Criterion and Google Benchmarks to benchmark our Rust functions simple_n2_cluster, simple_n2_cluster_simd, and tiled_n2_cluster against the original C++ versions of the code. We initially wanted to also compare our implementation against the Julia code as well but after discovering that we were actually faster than the C++ implementation we thought that deserved our full attention.
Previously, we just used the single-event.dat file located in one of the example data folders of Fastjet to benchmark all of our code. This was good for debugging since we could make sure that our final calculated jets were the same as Fastjet but for performance testing we wanted a more diverse dataset, with varying numbers of particles to process.
Luckily for us, there’s a way to generate simulated events using a package called Pythia. We spent some time generating a diverse dataset of events ranging from 50 to 1500+ particles per event and ran all 3 algorithms on it.
Let’s start with a macro view of our performance. Tiled clustering in C++ is still the fastest out of all the algorithms, however our tiling implementation is a pretty close second. The simple C++ implementation is also the slowest out of the algorithms with our SIMD and simple clustering consistently beating C++.
To us, this is pretty surprising! We assumed that Rust would at best match C++ performance with the added benefit of memory safety (which was already a huge win from our perspective).
Zooming in to a particular section gives us some more information too. Our tiled clustering algorithm has a slightly more unstable performance but the same points as the C++ code. We posit that this could be related to memory allocation and cache misses during reading tile neighbors? However, this would require some benchmarking with Memory profiling packages like dhat that would allow us to observe the exact heap allocations made during each run.
Another note here is that the benefits from a lot of the optimizations have a constant factor overhead (SIMD, tiling) so their benefit is really only felt at large input sizes. In fact, you can see that for smaller inputs, the simple n2 algorithm often outperforms our simd or tiled implementations.
We hope by now it’s clear that paying attention to how code gets compiled can really pay off compared to manually hand tuning for performance. Who knew all of these seemingly trivial changes could actually turn out to be so powerful? Just as important though, is the fact that we could actually measure and test the impact of these changes with Criterion.
Never forget to benchmark your code folks!
Until next time then :)