Exploring whether a Rust/SIMD backend could be useful to PyWavelets

[AlephNotation]

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.

[rgommers]

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:

• Are you planning to maintain this project/code long-term?
• Are you planning to keep adding features, or will it remain 1-D only?
• Can you run the same-interpreter benchmark on x86-64 with AVX2 and/or x86-64-v4?
• What are the wheel sizes per platform for a release build?
• Your results seem to be "planned" or "cold" - isn't the "planned + cold" worst-case?

[AlephNotation]

Taking these two as a pair

Are you planning to maintain this project/code long-term?
Are you planning to keep adding features, or will it remain 1-D only?

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.

Can you run the same-interpreter benchmark on x86-64 with AVX2 and/or x86-64-v4?

For sure. I'll go spin up some machines later this weekend since gh runners can be noisy for these kind of benchmarks.

What are the wheel sizes per platform for a release build?

I built a 25-wheel matrix using Rust 1.98, Maturin 1.14.1, and CPython 3.10–3.14:

• Linux x86-64: 439.6–442.8 KiB
• Linux aarch64: 385.1–390.3 KiB
• macOS x86-64: 391.7–393.3 KiB
• macOS arm64: 345.2–352.9 KiB
• Windows x86-64: 284.8–287.9 KiB

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.

Your results seem to be "planned" or "cold" - isn't the "planned + cold" worst-case?

Yeah these labels are confusing. I should fix them. They are alternative measurements, not additive:

The cold is already plan_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]

Thanks, that all sounds good. The binary size is reasonable - and I think it can probably be reduced on x86-64, because fearless_simd is 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.

I should probably rename these columns to 'reused plan' and 'plan + execute'.

That would have been clear to me indeed.

Ideally I would like the implementation priority to be guided by the PyWavelet maintainers.

I gave it some thought, and also had Claude do some ad-hoc benchmarking. cwt is 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 so axis=0 is much slower than axis=-1 (could be improved directly in this code base as well) - and batching/fusing may help.
  • That should improve performance for dwt2/dwtn and wavedec2/wavedecn. Relevant because used in scikit-image, which is probably the largest user of this package.
• Assuming the above works out: swt_axis - should be cheap increment once (1) exists.

[AlephNotation]

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, benchmark axis=0 and axis=-1 across 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]

I was officiating a wedding over the weekend and had much less time than I had hoped.

That sounds like more fun than writing code:)

---

Thanks for the update! All 4x - 5x, that looks like a consistent win across platforms.

I think the best approach would be to do a straightforward FIR implementation

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]

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]

The rust code is linked as a small separate binary with a tiny small C ABI.

That could mean cdylib or staticlib, which one do you have in mind? Either one is workable; staticlib (with no_std for 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=false to start with opt-in when building from source (setting it to true for wheel builds in CI), then flip the defaults once it's clear it's all stable.

[AlephNotation]

Already have it set up with staticlub. I'd like to push back against no_std for 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]

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.