Repository navigation
Exploring whether a Rust/SIMD backend could be useful to PyWavelets #865
Description
Hi everyone!
I started this project while exploring portable SIMD programming in Rust. Wavelet transforms seemed like a particularly good fit, and PyWavelets was unusually valuable as a source of ground truth.
After a few days I now have a safe Rust implementation of real-valued 1D DWT/IDWT and multilevel wavedec/waverec. I got curious as to whether we the SIMD implementation gains would hold through a python interface equivalent to PyWavelets.
On an apple M4 max, measured in the same CPython interpreter with the same numpy inputs:
- The reusable-plan api is 3.54× faster at the median across 70 cases and wins all 70.
- A conservative cold path, including wavelet construction and planning on every rust call, is 2.14× faster at the median and wins 64/70.
- For representative 4096 point f64 transforms, the cold-path gains range from 1.52× to 2.80×.
Output creation and destruction are included. The repository contains the methodology, compatibility tests, environment metadata, and all 4,200 raw timing samples:
- Results: https://github.com/AlephNotation/wavelets#python-to-python-api
- Raw samples: https://github.com/AlephNotation/wavelets/blob/master/benchmarks/results/apple-m4-max-python-api.json
- Python bindings: https://github.com/AlephNotation/wavelets/tree/master/python
I realize this is coming out of nowhere, and adding a rust toolchain would be a substantial deviation for PyWavelets. I want to be super clear that I am in no way suggesting that the existing implementation should simply be replaced. Besides, my current implementation only covers a subset of PyWavelets.
That said, if this is something you guys would be interested in, I would be interesting in doing the work.
If the build or maintenance implications make Rust a nonstarter, I completely understand. I wanted to ask before investing in a deeper integration prototype.
Activity
rgommers commentedon Aug 28, 2026
Hey @AlephNotation, thanks for sharing! Those gains are potentially interesting. A Rust toolchain is a heavy extra dependency, but I'm willing to explore it - at a minimum, we'll learn something in the process:)
Let me ask a few of what I think are the most critical questions:
AlephNotation commentedon Aug 28, 2026
Taking these two as a pair
Yeah that was my thought. It would be rude to dump a pile of code on you guys and walk away.
If the PyWavelets team is interested, I would work toward feature parity in agreed upon stages rather than leave the current 1-D implementation as some strange rust orphan. I've been playing around with paraunitary lattices and have figured out a way to make them SIMD efficient in a way that I think will scale even better in higher dimensions.
I would not want to touch PyWavelets current public API, numerical behavior, etc.
Ideally I would like the implementation priority to be guided by the PyWavelet maintainers. This is your project and I have no desire to unilaterally start rewriting things.
For sure. I'll go spin up some machines later this weekend since gh runners can be noisy for these kind of benchmarks.
I built a 25-wheel matrix using Rust 1.98, Maturin 1.14.1, and CPython 3.10–3.14:
These are compressed wheels. The whole thing is 9.06mb.
I can't yet do a true apples to apples comparison since PyWavelet has a much broader feature set.
Yeah these labels are confusing. I should fix them. They are alternative measurements, not additive:
The
coldis alreadyplan_factory().forward(...)(or inverse) inside one timed operation. So my 'cold' is already planning + execution. I think this current benchmark setup slightly favors PyWavelets, but if you have an alternative way of measuring I'm all ears!I should probably rename these columns to 'reused plan' and 'plan + execute'.
rgommers commentedon Aug 29, 2026
Thanks, that all sounds good. The binary size is reasonable - and I think it can probably be reduced on x86-64, because
fearless_simdis choosing four levels (see docs) and at least SSE 4.2 isn't useful anymore; if you care at all about performance, you're always going to have AVX2 at a minimum.That would have been clear to me indeed.
I gave it some thought, and also had Claude do some ad-hoc benchmarking.
cwtis popular but not a good candidate (uses FFT/convolve). I think the biggest wins may come from:dwt_axis/idwt_axis- looks like there's a lot to gain there, we're not doing even a transpose soaxis=0is much slower thanaxis=-1(could be improved directly in this code base as well) - and batching/fusing may help.dwt2/dwtnandwavedec2/wavedecn. Relevant because used inscikit-image, which is probably the largest user of this package.swt_axis- should be cheap increment once (1) exists.AlephNotation commentedon Sep 1, 2026
Sorry for the delay. I was officiating a wedding over the weekend and had much less time than I had hoped.
The x86 benchmarks are here. On the avx-2 machine reused-plan execution was about 4x faster at the median. On avx-512 the median was 4.82x.
The AVX-512 numbers are much stronger for long filters, but that is a partially a result of my SIMD lattice work. I think the lattice stuff is pretty interesting, but I want to be mindful of the mathematical complexity I would be introducing here.
I think the best approach would be to do a straightforward FIR implementation on
dwt_axis/idwt_axis, benchmarkaxis=0andaxis=-1across representative 2D/nD shapes. That would give us a ground truth baseline before we consider introducing more complicated math.I can start working on this in a fork.
rgommers commentedon Sep 2, 2026
That sounds like more fun than writing code:)
Thanks for the update! All 4x - 5x, that looks like a consistent win across platforms.
Agreed. The complicated math is best avoided, it would interfere with the "does Rust SIMD help, and can we introduce it" question. And it's more time-consuming to think about new math, and I don't have much time to spend on PyWavelets (and other maintainers have been inactive for the last couple of years).
AlephNotation commentedon Sep 2, 2026
I need to rerun benchmarks in the morning. AVX-2 is free of the fancy math and shows pretty good results, but the AVX-512 may still select the lattice backend for long filters. I'll rerun the AVX-512 in the morning without the lattices.
I'll try to keep these changes as small as possible given your bandwidth. The existing Cython _dwt extension continues to own the python and numpy facing behavior. The rust code is linked as a small separate binary with a tiny small C ABI. This way we can avoid introducing a separate module or changeing the public api.
I should have the for ready in a day or two. Don't worry about a quick review, we're all busy. I'm also happy to help with smaller maintenance work where useful.
rgommers commentedon Sep 2, 2026
That could mean
cdyliborstaticlib, which one do you have in mind? Either one is workable;staticlib(withno_stdfor size reduction) might be preferable. I know how to ship a shared library as well though (xref https://mesonbuild.com/meson-python/how-to-guides/shared-libraries.html).I'm thinking this should be integrated in the build directly, so Meson understands the build graph and builds the Rust code with
rustc. And probably a build option-Duse-rust=falseto start with opt-in when building from source (setting it totruefor wheel builds in CI), then flip the defaults once it's clear it's all stable.AlephNotation commentedon Sep 2, 2026
Already have it set up with
staticlub. I'd like to push back againstno_stdfor the time being. Very open to going that route, but I would like to settle the "does Rust SIMD help, and can we introduce it" question before I take on the additional engineering burden.Is that ok with you?
Question about Meson: Cargo support only landed in 1.11 and PyWavelets uses 1.1. Right now I have the 1.1 target supported through a opaque Cargo
custom_target. Are you ok with a version bump or do you prefer me to express the rust dep explicitly as a Meson subproject?rgommers commentedon Sep 2, 2026
Sure, that makes perfect sense - first things first, and perhaps it won't even be needed, depends on binary size of the end result.
A Meson version bump to >=1.11 is fine with me.