SIMD PAIR PROGRAMMING WITH THE COMPILER
In the process of working on Legato, I found a need to optimize sine wave generation in the project, and SIMD was an obvious candidate.
SIMD, or Single Instruction, Multiple Data, is a hardware feature that allows processors to perform operations on multiple pieces of data at the same time, using a single instruction. It's commonly used in audio and image processing.
As a side-note, also check out SIMT, the GPU cousin
Using SIMD intrinsics and libraries can provide incredible speedups, but it's not a tool that can be applied in every situation.
One thing you can find quite quickly is that the compiler is often much smarter than you, which is demonstrated quite wonderfully in this blog post here by Matklad.
Additionally, you may find yourself bottlenecked by CPU cache, allocations, system calls, or other costly operations, before the need for SIMD itself.
But there are certain situations where you might be able to outsmart the compiler, if you can spot certain patterns.
This is by no means a finished solution—likely there may be some further optimizations down the line—but I hope it can serve as an interesting lens for those that are interested in applying SIMD techniques to their own projects.
Some Credit is Due
I am grateful for the resources that enabled me to explore this domain, and I want to credit some brainstorming with a coworker, as well as the incredible Algorithmica resource, which helps you move from a lens of Leetcode and flaky benchmarks, to understanding cache aware data-structures, branchless programming, and many more performance optimizations.
I truly cannot believe this course is free. There is at least a master's degree worth of content sitting here for the taking, and it helped me move from a more fullstack skillset, to building intuition for systems performance.
Additionally, the blog post Lookup Tables Are Bad Approximators, by Jatin Chowdhury, helped me find a few final optimizations to push this a bit further.
Autovectorization is Your Friend
The compiler can autovectorize, or write SIMD itself given it has an understanding of the target platform.
When working on this, I had a real desire to "beat" the compiler, and to try and find a solution with SIMD that would outperform what the compiler is capable of.
This is really not the right perspective to apply. Rather, you need to think about in what ways you can structure your code to apply these autovectorization gains, and only when you reach these limitations does the need for SIMD really arise. There were a few times where I thought I gained enough of an instinct to start understanding when I needed SIMD, and a few times I was wrong.
Interpolated delay lines, for instance, which are really just ring-buffers with fractional, interpolated indexing, are a good example. I immediately went for SIMD, while I was carrying multiple ring-buffers for each audio channel, resulting in poorer CPU cache performance, and I could have used bit masking tricks to avoid extra math work. SIMD helped eventually, but I applied this too early, when minimizing CPU cache misses should have been the first lens I applied.
But I Still Want to Beat The Compiler
Fair enough, but the compiler really still is your partner, I promise. To beat the compiler, you have to come up with a trick that it can potentially not find or safely apply.
We will start with a simple naive sine wave:
pub fn naive_scalar(freq: f32, sr: f32, phase0: f32, out: &mut [f32]) {
let mut phase = phase0;
let inc = freq / sr;
for s in out.iter_mut() {
phase = (phase + inc).fract();
*s = (phase * TAU).sin();
}
}It works, and it's not bad. But I want to run a lot of sine waves in my project, maybe a ton of different voices in an FM synth, for instance.
But we can already start to coax out a few vectorization gains.
One easy example in Rust is to simply use the chunks_exact() API, which given a power of two value (assumption, maybe you need exact lane size) may be auto-vectorized for you, automatically, given you are building this target CPU. I found this to be particularly useful when working on Unsized collections, where the compiler did not have as much information about the underlying structure.
This already made a slight improvement, but it was still noise at this point.
The next trick I looked at was a prefix sum technique. This is a cumulative sum or scan of elements, in a series. If you want to read more about this, I would highly suggest digging into the linked article.
Even though our math is associative, the compiler is not allowed to reorder the floating point operations here, and therefore cannot find more optimizations for this phase scan.
Now, we can take this trick and apply it to our naive sine wave to compute our phase increment, like so:
// Note: Using Nightly SIMD here
/// LANES is a cfg parameter depending on your target CPU
pub type Vf32 = Simd<f32, LANES>;
#[inline(always)]
pub fn simd_scan(mut x: Vf32) -> Vf32 {
let mut offset = 1;
while offset < LANES {
x += shift_right(x, offset);
offset <<= 1;
}
x
}
#[inline(always)]
fn shift_right(x: Vf32, by: usize) -> Vf32 {
match by {
1 => x.shift_elements_right::<1>(0.0),
2 => x.shift_elements_right::<2>(0.0),
4 => x.shift_elements_right::<4>(0.0),
8 => x.shift_elements_right::<8>(0.0),
_ => unreachable!(),
}
}
pub fn simd_realsin(freq: f32, sr: f32, phase0: f32, out: &mut [f32]) {
let inc_scalar = freq / sr;
let inc = Vf32::splat(inc_scalar);
let mut base = phase0;
for chunk in out.chunks_exact_mut(LANES) {
let running = simd_scan(inc);
let phase = (Vf32::splat(base) + running).fract();
let sample = (phase * Vf32::splat(TAU)).sin();
chunk.copy_from_slice(sample.as_array());
base = phase.as_array()[LANES - 1];
}
}This brought us down from 14.49 µs to 10.74 µs, which is not a bad gain.
However, upon inspecting the assembly for the output, I noticed that the CPU I was using did not have a vectorized path for the sine call. This meant that the sin calls fell back to a scalar libm loop, and were also being packed/unpacked.
I then decided to look into alternatives, either LUT or similar, and I came across the blogpost by Jatin Chowdhury, Lookup Tables Are Bad Approximators.
I then used these polynomial approximations, and wrote a SIMD version, which was quite fast (even for the scalar path), but helped push the final version even faster.
Here is the "final" (not what is actually in Legato). It's a bit of a mess, but we effectively have a sine-wave approximation, and use this in conjunction with our previous prefix scan trick. The coefficients were taken from the aforementioned blog.
#[inline(always)]
fn fast_mod_mhalf_half(x: Vf32) -> Vf32 {
x - x.round()
}
/// Sine of `x` turns (1 turn = 2*pi rad), valid on [-0.5, 0.5].
#[inline(always)]
fn sin_turns_core<const ORDER: usize>(x: Vf32) -> Vf32 {
let x_sq = x * x;
let y = match ORDER {
3 => {
let x_1_3 = Vf32::splat(-24.694_19) + Vf32::splat(50.140_33) * x_sq;
x * x_1_3
}
5 => {
let x_3_5 = Vf32::splat(63.661_51) + Vf32::splat(-54.084_73) * x_sq;
let x_1_3_5 = Vf32::splat(-25.116_73) + x_3_5 * x_sq;
x * x_1_3_5
}
_ => {
let x_q = x_sq * x_sq;
let x_5_7 = Vf32::splat(-66.094_78) + Vf32::splat(32.026_8) * x_sq;
let x_1_3 = Vf32::splat(-25.132_366) + Vf32::splat(64.787_45) * x_sq;
let x_1_3_5_7 = x_1_3 + x_5_7 * x_q;
x * x_1_3_5_7
}
};
y * (x + Vf32::splat(0.5)) * (x - Vf32::splat(0.5))
}
#[inline(always)]
pub fn sin_turns<const ORDER: usize>(x: Vf32) -> Vf32 {
sin_turns_core::<ORDER>(fast_mod_mhalf_half(x))
}
pub fn simd_poly<const ORDER: usize>(freq: f32, sr: f32, phase0: f32, out: &mut [f32]) {
let inc = Vf32::splat(freq / sr);
let mut base = phase0;
for chunk in out.chunks_exact_mut(LANES) {
let running = simd_scan(inc);
let phase = Vf32::splat(base) + running;
let sample = sin_turns::<ORDER>(phase);
chunk.copy_from_slice(sample.as_array());
base = phase.as_array()[LANES - 1];
}
}This brings us down to 2.35-2.47µs, down from our original scalar 14.41µs. This is quite a nice performance gain, and you can see even better results on AVX256/512.
But What If We Gave Our Friend More Information?
Note: this section is much more CPU specific, and less reproducible. We are reaching microarchitecture optimizations, and your mileage may vary. This was on an M3 CPU.
Well, let's not give up on our friend yet. I mentioned earlier that the compiler was not allowed to reorder floating point operations? What if we could?
Well, Rust recently introduced a new feature to do so, somewhat of a local counterpart to the --ffast-math flag, algebraic operators.
pub fn poly_scalar_chunked_associative<const ORDER: usize>(
freq: f32,
sr: f32,
phase0: f32,
out: &mut [f32],
) {
let mut phase = phase0;
let inc = freq / sr;
for chunk in out.chunks_exact_mut(LANES) {
for s in chunk.iter_mut() {
phase = phase.algebraic_add(inc);
*s = sin_turns_scalar::<ORDER>(phase);
}
}
}If we instead continue with the auto vectorized version, and allow the compiler to reorder the math here, we suddenly find ourselves going even faster, at 1.6-1.79µs. The compiler here actually decided to just precompute various phase increments and apply them with a vectorized add (this optimization was lost with audio rate frequency, so it's only useful for LFO-like operators), and the polynomial math here was not using SIMD.
Curious to see if I could find any further optimizations, I wondered if breaking phase computing into a separate loop might help the compiler here, as we are removing a data dependency, and also give us a stronger locality of reference.
pub fn poly_fission_full<const ORDER: usize>(freq: f32, sr: f32, phase0: f32, out: &mut [f32]) {
let inc = freq / sr;
let mut phase = phase0;
for s in out.iter_mut() {
phase = phase.algebraic_add(inc);
*s = phase;
}
// dependency-free map, using algebraic poly
for s in out.iter_mut() {
*s = sin_turns_scalar_alg::<ORDER>(*s);
}
}This brings us down to a median runtime of 1.145 µs.
Lastly, you might notice here that we are actually iterating through a decent number of items. The M3 has a ludicrous amount of L1, but we can potentially keep these values hotter by chunking and computing them in smaller chunks. There may be a better term for this, but I mostly found terms surrounding bandwidth, roundtrip, etc. I am still learning here and would appreciate more insight. This blog here is relevant although may not be the actual mechanism occurring here.
pub fn poly_fission_tiled<const ORDER: usize, const TILE: usize>(
freq: f32,
sr: f32,
phase0: f32,
out: &mut [f32],
) {
let inc = freq / sr;
let mut phase = phase0;
let mut scratch = [0.0f32; TILE];
let mut chunks = out.chunks_exact_mut(TILE);
// Specifically work on certain chunk sizes first
for chunk in chunks.by_ref() {
// Precompute the phase
for slot in scratch.iter_mut() {
phase = phase.algebraic_add(inc);
*slot = phase;
}
// Compute the actual waveform
for (o, &p) in chunk.iter_mut().zip(scratch.iter()) {
*o = sin_turns_scalar_alg::<ORDER>(p);
}
}
// Process the tail
let tail = chunks.into_remainder();
for o in tail.iter_mut() {
phase = phase.algebraic_add(inc);
*o = sin_turns_scalar_alg::<ORDER>(phase);
}
}This brought us down to a final 968 ns runtime with a block size of 64.
Final Benchmarks
Here, we can see a few different benchmarks (Mac Air M3), for combinations of polynomial approximations, scalar, autovectorization vs SIMD. You will also see separate benchmarks for FM, as this data dependency changes how aggressively the compiler can operate.
I want to disclose that I used AI to create these various benchmark harnesses for this blog, but the actual original optimizations were not done by any AI.
You can find the repository here. Feel free to contribute any further optimizations.
Timer precision: 41 ns
fm fastest │ slowest │ median │ mean │ samples │ iters
├─ fm_fission_tiled_ │ │ │ │ │
│ ├─ 3 3.686 µs │ 4.249 µs │ 3.853 µs │ 3.839 µs │ 100 │ 200
│ ├─ 5 3.707 µs │ 4.416 µs │ 3.749 µs │ 3.778 µs │ 100 │ 100
│ ╰─ 7 3.874 µs │ 5.124 µs │ 3.936 µs │ 3.969 µs │ 100 │ 100
├─ fm_scalar_ │ │ │ │ │
│ ├─ 3 4.832 µs │ 6.791 µs │ 4.874 µs │ 4.976 µs │ 100 │ 100
│ ├─ 5 5.207 µs │ 9.374 µs │ 5.29 µs │ 5.682 µs │ 100 │ 100
│ ╰─ 7 5.915 µs │ 6.541 µs │ 5.999 µs │ 6.009 µs │ 100 │ 100
╰─ fm_simd_scan_ │ │ │ │ │
├─ 3 2.374 µs │ 3.29 µs │ 2.436 µs │ 2.462 µs │ 100 │ 200
├─ 5 2.478 µs │ 3.457 µs │ 2.519 µs │ 2.581 µs │ 100 │ 200
╰─ 7 2.603 µs │ 15.76 µs │ 2.645 µs │ 2.807 µs │ 100 │ 200
Timer precision: 41 ns
sine fastest │ slowest │ median │ mean │ samples │ iters
├─ autovec_attempt_ 13.66 µs │ 15.45 µs │ 13.79 µs │ 13.82 µs │ 100 │ 100
├─ naive_scalar_ 14.12 µs │ 19.79 µs │ 14.49 µs │ 14.71 µs │ 100 │ 100
├─ poly_fission_full_ 1.124 µs │ 1.645 µs │ 1.145 µs │ 1.193 µs │ 100 │ 400
├─ simd_realsin_ 10.33 µs │ 24.74 µs │ 10.83 µs │ 11.63 µs │ 100 │ 100
├─ poly_fission_tiled_ │ │ │ │ │
│ ├─ 16 1.395 µs │ 2.03 µs │ 1.416 µs │ 1.456 µs │ 100 │ 400
│ ├─ 32 978.5 ns │ 2.655 µs │ 1.02 µs │ 1.087 µs │ 100 │ 400
│ ├─ 64 1.009 µs │ 1.468 µs │ 1.03 µs │ 1.065 µs │ 100 │ 400
│ ├─ 128 947 ns │ 1.426 µs │ 968 ns │ 1.017 µs │ 100 │ 400
│ ╰─ 256 1.009 µs │ 1.499 µs │ 1.02 µs │ 1.09 µs │ 100 │ 400
├─ poly_scalar_ │ │ │ │ │
│ ├─ 3 4.54 µs │ 6.624 µs │ 4.624 µs │ 4.768 µs │ 100 │ 100
│ ├─ 5 4.915 µs │ 7.416 µs │ 4.957 µs │ 5.191 µs │ 100 │ 100
│ ╰─ 7 5.624 µs │ 8.29 µs │ 5.666 µs │ 5.882 µs │ 100 │ 100
├─ poly_scalar_chunked_ │ │ │ │ │
│ ├─ 3 3.79 µs │ 5.499 µs │ 3.874 µs │ 4.228 µs │ 100 │ 100
│ ├─ 5 3.874 µs │ 12.37 µs │ 4.082 µs │ 4.34 µs │ 100 │ 100
│ ╰─ 7 3.999 µs │ 12.95 µs │ 4.166 µs │ 4.652 µs │ 100 │ 100
├─ poly_scalar_chunked_associative_ │ │ │ │ │
│ ├─ 3 1.582 µs │ 21.43 µs │ 1.645 µs │ 2.088 µs │ 100 │ 200
│ ├─ 5 1.686 µs │ 2.499 µs │ 1.728 µs │ 1.785 µs │ 100 │ 200
│ ╰─ 7 1.728 µs │ 2.665 µs │ 1.79 µs │ 1.845 µs │ 100 │ 200
╰─ simd_poly_ │ │ │ │ │
├─ 3 2.207 µs │ 3.27 µs │ 2.228 µs │ 2.352 µs │ 100 │ 200
├─ 5 2.249 µs │ 3.436 µs │ 2.291 µs │ 2.515 µs │ 100 │ 200
╰─ 7 2.374 µs │ 3.52 µs │ 2.415 µs │ 2.522 µs │ 100 │ 200These numbers are medians from a laptop and should be read as a relative overlook, there is some noise here
Summary
We never actually "beat" the compiler. This is not to say that SIMD is slower and that the compiler is always faster, but I will make the argument that the compiler was able to find some clever optimizations that I had not considered. Legato uses the poly_simd implementation currently, as the waveforms can take in an audio rate frequency stream, but the compiler did find an approach much faster than mine, without the data dependency.
I will argue that writing SIMD, for most developers, should be a sort of pair programming exercise with your compiler. You may try one technique, and find that by giving it the context or permission, the compiler can help you with other optimizations. By inspecting the generated assembly, you may find an area where the compiler was unable to apply vectorization, or you may take some inspiration from various techniques it found.
The compiler is your friend, and it really enjoys when you give it the same context you have. Don't try to beat it, work with it back and forth until you both run out of ideas.