ACP: Add functions for "compress bits" and "expand bits" to unsigned integers

[okaneco]

Proposal

Problem statement

There's currently no way to take advantage of compiler optimizations or hardware instructions if a naive implementation is written to extract multiple bits from an integer or deposit those bits into an integer according to a mask.

These operations (known as parallel bit extract/deposit, gather/scatter, or compress/expand) are useful building blocks for algorithms and data structures, and they are backed by hardware instructions on some platforms.

I will refer to these functions as some variant of "bit expand/compress" since I find those names more descriptive.

Motivating examples or use cases

Bit expand and compress would be good candidates for std inclusion because they are non-trivial to implement efficiently and may require the use of special intrinsics from the compiler to optimize portably according to the target platform, like the other bit operations of trailing_zeros and count_ones.

Expand and compress have use cases for working with variable length encodings (varints, Unicode), bit interleaving/deinterleaving (Morton codes/space-filling curves), perfect hashing/lookup tables, parsing tasks, and succinct data structures. They're also useful for expressing operations that would ordinarily take several masks and shifts, and simplifying that to one function call with a specified mask.

I'm not aware of any pattern that will get optimized into the desired instruction on x86 so the intrinsics must be used.
LLVM doesn't currently have a generic intrinsic but there has been discussion on the Discourse and there's a C++ proposal for C++26 that would include the functionality so it may be implemented in the future.

Solution sketch

A naive implementation is shown for each function. Nonzero and signed integers would be omitted.

/// Returns `self` with the bit locations specified by `mask` packed
/// contiguously into the least significant bits of the result.
/// ```
/// let n = 0b1011_1100;
///
/// assert_eq!(compress_bits(n, 0b0010_0100), 0b0000_0011);
/// assert_eq!(compress_bits(n, 0xF0), 0b0000_1011);
/// ```
pub const fn compress_bits(self_: u32, mask: u32) -> u32 {
    let mut result = 0;
    let mut i = 0;
    let mut j = 0;
    while i < u32::BITS {
        let mask_bit = (mask >> i) & 1;
        result |= (mask_bit & (self_ >> i)) << j;
        j += mask_bit;
        i += 1;
    }
    result
}

/// Returns a result with the least significant bits of `self` distributed to
/// the bit locations specified by `mask`.
/// ```
/// let n = 0b1010_1101;
///
/// assert_eq!(expand_bits(n, 0b0101_0101), 0b0101_0001);
/// assert_eq!(expand_bits(n, 0xF0), 0b1101_0000);
/// ```
pub const fn expand_bits(self_: u32, mask: u32) -> u32 {
    let mut result = 0;
    let mut i = 0;
    let mut j = 0;
    while i < u32::BITS {
        let mask_bit = (mask >> i) & 1;
        result |= (mask_bit & (self_ >> j)) << i;
        j += mask_bit;
        i += 1;
    }
    result
}

A more optimal way to implement this is with an XOR parallel suffix as discussed in Hacker's Delight, or with native instructions.
The function can optimize to only a few instructions when the mask is known, as shown here.

Alternatives

1. Do nothing and rely on third party crates implementing this in software or wrappers around the platform intrinsics.
2. Wait for LLVM to implement the intrinsic.

As for naming, I think suitable Rust names would be either of the first two.
"Extract" is in my opinion too overloaded with meaning to be a good name for what it does.

• compress_bits, expand_bits
• gather_bits, scatter_bits
• extract_bits, deposit_bits

Links and related work

PEXT, PDEP - Intel BMI2 (introduced with Haswell BMI2)
BEXT, BDEP - Arm SVE2
APL's Replicate and Expand
https://orlp.net/blog/extracting-depositing-bits/
https://en.wikipedia.org/wiki/X86_Bit_manipulation_instruction_set#Parallel_bit_deposit_and_extract
https://eisenwave.github.io/cpp-proposals/bit-permutations.html#interleaving-bits - C++26 proposal examples
https://discourse.llvm.org/t/how-would-i-go-about-implementing-new-bit-manipulation-builtins-for-my-proposal-especially-generic-builtins/76523

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[hanna-kruppe]

I would definitely recommend waiting for an LLVM intrinsic. If LLVM gains a lot of cleverness to synthesize the best ways to implement this operation with various fixed masks on different targets, it would be useful to expose that. But in the current state of the world, especially with runtime masks, I believe this function will be a serious performance footgun:

• A generic, portable software fallback with variable mask is extremely inefficient. For most use cases, the actual problem is massaged into an odd shape to be able to use the efficient hardware instructions. If the hardware instructions aren't available, emulating the compress/expand is often much worse than solving the problem at hand directly, using whatever was the most efficient option before PEXT/PDEP came along.
• The relevant instructions are relatively recent and not available in the baseline targets most people use. Thus, using the relevant instructions at all will often require the user to do runtime target_feature detection and function multi-versioning. Even if the standard library could do the runtime feature detection in compress_bits and expand_bits (not yet possible in core), for efficiency the multi-versioning has to happen at a higher level, not on every single compress/expand.
• There are more efficient fallbacks for some targets but these, too, require target features not included in most baselines. So even if you have some code that really needs to emulate the full generality, you'll want some (runtime) feature detection to get a decent fallback. In fact, it's not just target_feature detection. Since some AMD CPUs before Zen3 had infamously slow implementations, some software used "has PEXT/PDEP and isn't Zen1 or Zen2" as condition for selecting which code to use.
• Availability of these instructions is still effectively limited to x86 (edit: and powerpc, thanks @programmerjake), making them a performance portability trap. The SVE2 instructions mentioned are only for vectors, not for scalars. There's a risk that people will use this in nominally portable code but only test performance on their x86 CPU with -Ctarget-cpu=native or similar (maybe set in $CARGO_HOME/config.toml long ago).

[programmerjake]

• Similarly, Power10 has vec_pext and vec_pdep.

it also has scalar equivalents.

[okaneco]

But in the current state of the world, especially with runtime masks, I believe this function will be a serious performance footgun

I also considered initially only exposing constant mask functions and waiting for the LLVM support for variable masks. This would limit performance missteps.

let bits = x.compress_bits::<0b1010_0011>();

pub const fn compress_bits<const MASK: u32>(&self) -> u32 { self.compress_bits_var_mask(MASK) }
const fn compress_bits_var_mask(&self, m: u32) -> u32 { ... }

I would definitely recommend waiting for an LLVM intrinsic. If LLVM gains a lot of cleverness to synthesize the best ways to implement this operation with various fixed masks on different targets, it would be useful to expose that.

I'm not sure if this is what you meant, but LLVM is able to generate fixed masks if the complex (and somewhat difficult to understand) implementation is used, which should optimize portably. This is the method I alluded to in the solution sketch.

https://rust.godbolt.org/z/fMaza7vE8 - compress_bits with const mask

When writing the straightforward loop form like the right, the known masks still don't get optimized as well depending on the bit pattern. It sees through some simple shifts but falls over on others.

There are more efficient fallbacks for some targets but these, too, require target features not included in most baselines.

count_ones is similar in this regard. popcnt is not part of the x86/x86_64 baseline but its functionality is still useful in the fallback form. I think the functionality of the bit functions in this ACP would still be an improvement over not having them in the standard library, but I understand if it's decided to wait.

[hanna-kruppe]

I also considered initially only exposing constant mask functions and waiting for the LLVM support for variable masks. This would limit performance missteps.

Yes, but the const-mask-only API would become completely redundant if the variable-mask variant is added later, so stabilizing it is a tough sell.

I'm not sure if this is what you meant, but LLVM is able to generate fixed masks if the complex (and somewhat difficult to understand) implementation is used, which should optimize portably. This is the method I alluded to in the solution sketch.

As you're probably aware, there's cases where one can do better than shift/and/or with fixed masks (e.g., using multiplication). I would expect optimizations along those lines to only be applied reliably(!) if LLVM adds an intrinsic for it.

count_ones is similar in this regard. popcnt is not part of the x86/x86_64 baseline but its functionality is still useful in the fallback form. I think the functionality of the bit functions in this ACP would still be an improvement over not having them in the standard library, but I understand if it's decided to wait.

Popcount is similar in that it's not available on a popular target's baseline feature set (x86-64) but there are also significant differences:

• I'm not aware of software fallbacks for popcount that benefit significantly from runtime feature detection. Typical fallback implementations require only basic integer operations that are available unconditionally (except for those few poor souls targeting RISC-V without the M extension).
• Popcount is also far more widely available on other architectures, e.g., AAch64 and RVA22 require it, and even on x86 "has popcount" is a significantly larger set than "has fast PEXT/PDEP" despite everything being in the BMI2 extension.
• Software implementation are slower than hardware popcount, but still fast enough to be useful for many things -- that's how we got the rich literature on software implementations.
• As far as I'm aware, there aren't as many cases as with PEXT/PDEP where a software fallback for popcount can be made significantly faster by solving a simpler problem directly.

At risk of over-simplifying: popcount instructions were added because lots of software needs this specific operation, while PDEP/PEXT originates from hardware people realizing that they can efficiently implement a pair of operations that are applicable for lots of different things various software is doing. Now that the instructions exist in some hardware, people also come up with software designs that are only attractive when these operations are efficient, but that's still in its infancy compared to popcount.

[Amanieu]

We discussed this in the @rust-lang/libs-api meeting. We don't think that these should be stabilized before explicit LLVM exists for these intrinsics. However it would be fine to add these as unstable for now for experimentation. In the discussion there was a preference for gather_bits/scatter_bits.

Please open a tracking issue and open a PR to rust-lang/rust to add it as an unstable feature. You can close this ACP once the tracking issue has been created.

[programmerjake]

I'll note I prefer compress/expand over gather/scatter because by analogy to SIMD, SIMD gather/scatter allows you to arbitrarily reorder elements (it is a fully-general load/store to independently-determined addresses for each element), whereas SIMD compress/expand does not allow reordering the loaded/stored elements, just inserting/removing elements.

[okaneco]

I've created the tracking issue, noting the criteria of having explicit support for this in LLVM before stabilization

rust-lang/rust#149069