Skip to content

blosc: AVX2 shuffle is compiled out of x86_64 wheels, and aarch64 gets the scalar shuffle #181

Description

@kylebarron

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:

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:

$ 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:

#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:

// 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. 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 — https://github.com/mulimoen/rust-blosc-src
  • 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions