developmentseed / developmentseed/zarrista
blosc: AVX2 shuffle is compiled out of x86_64 wheels, and aarch64 gets the scalar shuffle
- Dominant language
- Python
- Stars
- 45
- Forks
- 2
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 7
Description
Our blosc codec runs a slower shuffle implementation than the hardware supports, on every platform we ship. The runtime CPU detection in c-blosc works correctly. The faster implementations are simply not compiled into the binary for it to select.
This is a build-configuration problem in the `blosc-src` crate, not a bug in zarrs or zarrista. We can fix it for x86_64 with a small patch. We cannot fix it for aarch64 at this layer.
## Background
Blosc applies a byte transposition (shuffle, or bitshuffle at bit granularity) to each block before it calls the sub-codec. This groups bytes of equal significance together, which is where most of blosc's compression ratio comes from. The shuffle touches every byte of every chunk, in both directions, so it is on the hot path for all reads and writes.
c-blosc ships three implementations of this routine: generic (scalar), SSE2, and AVX2. It selects between them in two layers:
1. **Compile time.** The `SHUFFLE_SSE2_ENABLED` and `SHUFFLE_AVX2_ENABLED` defines decide which implementations are compiled into the library at all.
2. **Run time.** `blosc_get_cpu_features()` issues a real `cpuid`, confirms that the OS has enabled YMM state through `xgetbv`, and then `get_shuffle_implementation()` returns a function-pointer table for the best available implementation.
The runtime layer is careful and correct. It can only choose among implementations that the compile-time layer admitted.
## Problem 1: x86_64 wheels never get the AVX2 shuffle
`blosc-src`'s `build.rs` decides the compile-time layer from `CARGO_CFG_TARGET_FEATURE`, which describes the features the *target* is assumed to have — not the features of the machine that runs the wheel:
```rust
if target_features.contains(&"avx2") {
build.define("SHUFFLE_AVX2_ENABLED", "1");
build.flag("-mavx2");
}
```
A portable wheel must run on any x86_64 CPU, so we build for the baseline target, which does not include AVX2:
```console
$ rustc --print cfg --target x86_64-unknown-linux-gnu | grep target_feature
target_feature="fxsr"
target_feature="sse"
target_feature="sse2"
```
`SHUFFLE_AVX2_ENABLED` is therefore never defined, and the AVX2 branch is preprocessed out of the dispatcher in `shuffle.c`:
```c
#if defined(SHUFFLE_AVX2_ENABLED)
if (cpu_features & BLOSC_HAVE_AVX2) { ... return impl_avx2; }
#endif
```
Every x86_64 wheel we publish therefore tops out at the SSE2 shuffle, on hardware that has had AVX2 since 2013.
### Why `blosc-src` is written this way
`blosc-src` compiles all of c-blosc as a single `cc::Build`. Passing `-mavx2` would apply it to `blosc.c`, `blosclz.c`, `fastcopy.c`, and everything else, which would let the compiler auto-vectorize the whole library with AVX2 and crash with `SIGILL` on older CPUs. Gating the whole thing is a safe but blunt way to avoid that.
c-blosc's own CMake build avoids the problem by setting `COMPILE_FLAGS` on `shuffle-avx2.c` and `bitshuffle-avx2.c` alone. That is what makes its runtime dispatch safe.
## Problem 2: aarch64 wheels get the scalar shuffle
`blosc-src` vendors c-blosc **1.x**. Its `c-blosc/blosc/` directory contains only `shuffle-generic.c`, `shuffle-sse2.c`, and `shuffle-avx2.c` (and the bitshuffle equivalents). There is no `shuffle-neon.c`, and `build.rs` has no aarch64 branch.
On ARM, compilation therefore falls into c-blosc's fallback arm — the one that emits `#warning Hardware-acceleration detection not implemented for the target architecture` — where `blosc_get_cpu_features()` is hardcoded to return `BLOSC_HAVE_NOTHING`. Our aarch64 wheels, and Apple Silicon development machines, always run the scalar shuffle.
c-blosc2 has a NEON implementation. c-blosc1 never got one.
## Impact
The cost depends on the sub-codec. With `blosclz` or `lz4` the shuffle is a large share of total codec time. With `zstd` at higher levels the compressor dominates and the shuffle matters less.
This matters for release benchmarking. An Apple Silicon machine measures the scalar shuffle, and x86 CI measures the SSE2 shuffle. Neither reflects what blosc can do, and the two platforms are handicapped by different amounts. An ARM-vs-x86 blosc comparison would partly measure this issue rather than the library. Any published blosc numbers should either come after the x86 fix, or carry this caveat.
## Proposed fix for x86_64
Split the single C build in `blosc-src`'s `build.rs` in two. Compile everything at baseline as today, but put the two AVX2 sources in their own compilation unit with `-mavx2`, and always enable the dispatch branch on x86:
```rust
// Everything except the AVX2 sources — baseline, safe on any x86_64.
let mut build = cc::Build::new();
add_files_except(&mut build, "c-blosc/blosc", &["shuffle-avx2.c", "bitshuffle-avx2.c"]);
build.define("SHUFFLE_AVX2_ENABLED", "1"); // enable the runtime-dispatch branch
build.compile("blosc");
// Only the AVX2 sources get -mavx2. Reached solely through the cpuid check.
let mut avx2 = cc::Build::new();
avx2.file("c-blosc/blosc/shuffle-avx2.c")
.file("c-blosc/blosc/bitshuffle-avx2.c")
.define("SHUFFLE_AVX2_ENABLED", "1")
.flag("-mavx2"); // "/arch:AVX2" on MSVC
avx2.compile("blosc_avx2");
```
The AVX2 routines are only ever called through the function pointers that `get_shuffle_implementation()` hands out after the `cpuid` check passes. No other translation unit can drift into AVX2 instructions. The result is one wheel that runs everywhere it does today and uses AVX2 when the CPU has it. There is no portability cost, unlike `-C target-cpu`.
### Delivery
1. Fork `blosc-src` and apply the patch.
2. Add a `[patch.crates-io]` entry for `blosc-src`, alongside the existing entries for `zarrs`, `monostate`, and `dlpark`. This works for transitive dependencies, so zarrs needs no change.
3. Send the patch to [mulimoen/rust-blosc-src](https://github.com/mulimoen/rust-blosc-src). Version 0.3.8 is the latest release and has no feature flags for this.
4. Remove the fork and the patch entry when the change is released.
This does not affect the emscripten wheel. zarrs uses the pure-Rust `blusc` crate on `wasm32`.
### Verification
Verification needs an x86_64 machine with AVX2. GitHub's `ubuntu-latest` runners qualify.
- Run `nm` on the built extension and confirm `blosc_internal_shuffle_avx2` is present. This proves the implementation is compiled in.
- Compare timings on a blosc-compressed array before and after. This proves it is selected.
## Options for aarch64
No patch to `blosc-src` can help. c-blosc 1.x has no ARM shuffle code to enable.
- Ask the zarrs maintainers about c-blosc2 (which has NEON) for native builds, or about whether the pure-Rust `blusc` crate that zarrs already uses on `wasm32` is competitive enough to use on ARM as well. This is a zarrs decision.
- Otherwise, accept the scalar shuffle for now and document the caveat on any published ARM benchmark.
## Rejected alternatives
- **`-C target-cpu=x86-64-v3` (or `target-feature=+avx2`) on release builds.** There is no wheel tag for a microarchitecture level, so a wheel built this way would fail with `SIGILL` on any CPU without AVX2. The only distribution-safe form is a second package, as Polars does with `polars-lts-cpu`. That is far more machinery than this warrants, and the fix above achieves the same result with no portability cost.
- **Setting `CFLAGS` to force the define and the flag.** `cc` applies `CFLAGS` to every crate in the build that compiles C, so this would also hit `zstd-sys` and `libz-sys`, and would apply `-mavx2` to all of c-blosc rather than to the two dispatched files. That is the `SIGILL` scenario with a wider blast radius.
- **Overriding the dependency's build script from our manifest.** Cargo has no such mechanism. A fork is the only route.
## References
- `blosc-src` 0.3.8 —
- Pulled in by `zarrs` with the `snappy`, `lz4`, `zlib`, and `zstd` features.
- `BLOSC_PRINT_SHUFFLE_ACCEL=1` makes c-blosc print the CPU features it detects. That code is inside the x86-only path, so it prints nothing on ARM.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with blosc-src's build.rs and the c-blosc/blosc/shuffle-avx2.c and bitshuffle-avx2.c sources to understand the current single-build configuration. Build the extension on an AVX2 x86_64 runner, use nm to confirm blosc_internal_shuffle_avx2 is present, and compare blosc timings before and after; done means baseline portability remains intact while AVX2 dispatch is available.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c, python, rust
- Domain
- build-system, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 56/100