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.