Part 1: Rewriting C++ code to Rust (a FastJet experiment)
particle physics x rust
what we are doing
TLDR: This blog documents what we learned rewriting a scoped down fastjet 3.5.1 from C++ to Rust as our first real Rust crate, fastjet-rs. We benchmark carefully hand-tuned work of physicists from 20 years ago and a more recent Julia rewrite against our Rust implementation. This project is largely educational, but we also wanted to produce a crate that could become a meaningful contribution to the Rust ecosystem.
One emerging area of Rust development is scientific high-performance computing. Workloads in computational biology or particle physics routinely handle petabytes of data and demand heavy optimization and parallelization to run. For this project, we took a particular interest in particle physics Rust crates since my collaborator started programming by analyzing lepton jets for CMS at the LHC (the world’s largest particle accelerator in CERN).
Scientific computing seemed like natural Rust territory, but some deep diving into reddit threads revealed that adoption outside of ML and computational biology is still nascent. In fact, Fortran and C/C++ still have a chokehold on HPC programs even when Rust is beginning to define modern programming elsewhere. The discussion boils down to three points:
- Scientists aren’t programers and rust has a steep learning curve
- C/C++ code when written well is already performant
- C/C++ ecosystems have decades of niche library support that Rust can’t match yet
Any memory bugs within applications like fastjet have been weeded out and corrected by the community over the past 20 years at the expense of many segfaults and memory leaks. When the community is already comfortable with a battle-tested Fortran or C program, fighting the borrow checker won’t always be worth it.
That said, particle physicists haven’t been standing still. In the past few years, Python has gained serious traction among scientists who value ease of use and development speed. Binding tools like pybind11, Cython, SWIG give user-friendly interfaces to the underlying C++ code. There is no doubt that this is a step in the right direction. But accumulating tech debt by layering abstractions over aging C++ code risks the same fate as Fortran, where supporting new features and maintenance becomes more similar to witchcraft than engineering. One new line of code breaks a thousand others without rhyme or reason.
Julia has emerged as one solution to this problem. Our hypothesis is that Rust could be another. Rust consistently matches or beats optimized C in benchmarks and its ownership model gives you fearless concurrency without a garbage collector. These are guarantees that Julia, for all its strengths, doesn’t offer at compile time.
Our goal is to test this hypothesis by benchmarking fastjet across C++, Rust, and Julia.
some basic things to know about jets and the problem we’re solving
I had no physics background coming into this project, so we started with a crash course.
Watching the two part series is not entirely necessary, although very entertaining. The production quality is peak 2006 internet and honestly that only makes them better.
Ok so with that, what are jets?
proton proton collision that produces quarks and gluons which get recorded by the detectors calorimeters as jets
As Matt Strassler’s explains, jets are sprays of high-energy subatomic particles (hence the name jets) composing the building blocks of the universe that we examine to understand the inner workings of protons, neutrons and other hadrons. The particles that produce jets are extremely short-lived and never reach our detectors, so jets are the only observable evidence of their existence.
Before we can analyze jets we need to reconstruct them from scattered event outputs collected by the calorimeters inside particle colliders. Historically, there are two methods for reconstructing particle collision events into jets: cone algorithms and hierarchical clustering. From the fastjet manual:
Cone algorithms put together particles within specific conical angular regions, notably such that the momentum sum of the particles contained in a given cone coincides with the cone axis (a “stable cone”).
What the manual calls sequential recombination algorithms can be understood more broadly as hierarchical clustering. Pairs of particles bound by their distance are clustered repeatedly until all particles are merged into stable jets. This statistical pattern can be found everywhere today. For our purposes, we can approximate jet reconstruction as an application of the nearest-neighbor problem common in CS.
Both methods rely on a tunable distance parameter delta R to define either the radius of a cone from the origin particle or the furthest distance between two points that we would collimate into the same jet. Set delta R too low and we miss lower energy descendant particles (normally called “children”) or too high and background noise from other particles are included in our final jets.
what did we do?
Fastjet’s has three core classes that configure how jets are reconstructed:
fastjet::PseudoJetwhich defines jet objects as four-momenta vectors with additional values needed for different clustering algorithmsfastjet::JetDefinitionwhich contains the specifications for how jet cluster should be carried outfastjet::ClusterSequencewhich is the class that outputs the final jets after performing clustering as specified by JetDefinition on input PseudoJet objects
It is important to us that we keep the functionality of fastjet-rs close to fastjet which meant reimplementing PseudoJet, JetDefinition, ClusterSequence code as best we could in a rust-y manner.
Heavy pointer manipulation and use of templates were the two characteristics of fastjet that made it difficult to translate to rust. You can read more about why raw pointer handling is an unsafe action in rust here, but the tldr is that dereferencing raw pointers breaks rust’s memory safety guarantees. Of course, the advantage of pointer manipulation when done correctly lies in the performance benefits. So the main challenge is really just being explicit with memory safety while maintaining performance parity with fastjet’s unencumbered C++ implementations.
Writing getters for PseudoJet
Each PseudoJet object is initialized by a four-momentum vector describing a particle’s x, y, z planar coordinates and an energy value e. These values are then used to derive the information we will need for clustering: kt2, rap, and phi. These values are expensive to calculate, and some jets may never need to call them so Fastjet initializes them lazily with an internal function _finish_init(). _finish_init()
Our first design of some of the PseudoJets get methods made them mutable. This seems fine at a glance but quickly complicated our code quite dramatically. During the clustering step we often need to calculate the distances between two jets using the following function:
fn _bj_dist<J: ProxyJet>(jet_a: &mut J, jet_b: &mut J) -> f64 {
let dphi: f64 = PI-f64::abs(PI-f64::abs(jet_a.phi()-jet_b.phi()));
let deta: f64 = jet_a.eta()-jet_b.eta();
dphi * dphi + deta * deta
}Since we’re only accessing fields of jet_a and jet_b, theoretically we don’t actually need to pass in mutable references to jet_a or jet_b. However, since phi is a lazily initialized value with a mutable get method this aspect is passed onto the _bj_dist function. Furthermore, jets are stored in a vector so when calling _bj_dist we need to pass in two mutable references to elements of the same vector. This creates a major violation in rust since there cannot be multiple, simultaneous mutable references to the same object.
Even though we could guarantee that there would be no data races when running our code single threaded, the Rust compiler does not. Our temporary fix was to create a helper function which breaks our jet vector into two separate mutable components each time we had to call _bj_dist. You can probably see that this was a messy solution, not to mention a big performance hit to our naive implementation that we needed to fix.
Our second solution was to turn to unsafe Rust, since as stated earlier we guaranteed that we would never be calling mut on the same index, we could reduce the bloat of our function by calling the pointers directly. Unsafe code isn’t always bad, Rust just forces you to be deliberate about the abstractions that often lead to nasty memory bugs.
Luckily, there is a much better way to solve this problem in Rust with a concept called interior mutability, a design pattern that allows you to mutate data even when there are immutable references to that data. I told you earlier that this is explicitly disallowed by rust, nothing has changed. Implementations of interior mutability like OnceCell are just safe APIs for unsafe code, making it easier for us to directly use helpful abstractions for unsafe code.
For lazy initialization specifically, the standard rust library has a type called OnceCell which we can use to obtain shared immutable references to uninitialized data that becomes initialized once written. Getting rid of our vector split workarounds improves the performance, readability of our code, and prepares us for properly implementing concurrency.
In fact, the authors of fastjet ran into this very problem with these getter methods when multithreading their code. There are many preprocessing macros modifying how these getters work for the exact reason that they can’t prevent data races here. This is where the rust compiler’s foresight really pays off, by designing early with restrictive borrowing and ownership rules from the start we get “fearless concurrency” for free!
benchmarking
To compare our implementations, we timed and generated flamegraphs for the C++, Julia, and Rust versions of the simple n² algorithm using the basic clustering strategy.
Setting up profiling turned out to be more annoying than expected, especially on mac. We found that sample worked best for our purposes, but getting useful flamegraphs required some care. We had to reduce inlining (-g for C++, #[inline(never)] in Rust) and keep frame pointers intact so the profiler could actually record meaningful stack frames instead of collapsing everything into “the clustering function”.
C++ Flamegraph
Rust Flamegraph
In our Rust flamegraph, we found that roughly half of the runtime was spent actual clustering while most of the remainder went to parsing event files. Some nice improvements can be seen when compared to earlier flamegraphs, calculations for bj_dist (which required calculating eta and phi) no longer take half the time that clustering currently takes.
Julia Flamegraph
This is where things got interesting. We initially thought Julia was running significantly faster than our Rust implementation, but through observing the flamegraph it turned out the Julia benchmarks were only measuring jet creation not event parsing or output. Once we accounted for that, the playing field leveled out considerably:
- Rust: 498.003 µs ± 0.687 µs
- Julia: 487.902 µs ± 22.156 µs
- C++: 365.222 µs ± 0.901 µs
Julia and Rust land in essentially the same ballpark, with C++ still comfortably ahead. Currently, Julia has similar performance to our Rust implementation because of SIMD (Single Instruction Multiple Data), which tells the compiler to vectorize loops in assembly. Rust currently only supports SIMD for floats in the nightly crate, so for this first iteration we wanted to avoid its use, with it we most likely will pull fairly ahead of Julia.
However since Julia is a runtime language, the first iteration of the code runs extremely slow, otherwise known as the ‘Time to First Plot” problem. Lengthened compile time in Rust is a deliberate sacrifice for faster runtime programs and guaranteed memory safety. Compared to C++, Rust loses on the compile time game, taking around 30 seconds to fully compile versus C++’s three seconds. Rust is thus a happy middle ground between Julia and C++ that balances speed and memory management.
So why is our Rust implementation slower than C++?
A few reasons most likely. Because we use indices to track nearest neighbors, we need an extra pass through all active particles after each merge to make sure NN pointers stay valid. That introduces an additional constant-factor loop inside the O(n²) main loop. One potential fix is wrapping each PseudoJet in an Rc, which would allow multiple references to mutable jet instances and eliminate the index bookkeeping. Although, there is a possibility that the reference counting overhead could end up costing more than the loop it replaces. We’re also paying a cost by passing entire vectors by reference into one of our distance functions, and we’re investigating whether leveraging const more aggressively in our Rust code could close the gap, since the C++ implementation makes heavy use of inline const templating.
That said, profiling and iterating paid off. With the help of the Rust Performance Book and some build configuration changes, we brought our runtime down from 1.4 milliseconds to 650 microseconds which is roughly a 2x improvement from relatively simple optimizations (lazy initialization with OnceCell and updating NN indices for only active particles rather than the full list.)!
next steps
Fastjet-rs is very much an ongoing project and there are many other features that still need to be implemented. To start, we want to increase the number of jet finding algorithms we support as we currently only have Anti-kt. Another major feature we are missing is an O(nlogn) Voronio clustering algorithm that leverages a fascinating geometry based optimization trick.
So if you’re interested in seeing what we learn from any of these small projects stay tuned for the next part :)