Keep only the 10 useful bits in 16-bit words

Question

I have _m256i vectors that contain 10-bit words inside 16-bit integers (so 16*16-bit containing only 16*10 useful bits). What is the best/fastest way to extract only those 10-bits and pack them to produce an output bitstream of 10-bit values?

Answer 1

Here’s my attempt.

Have not benchmarked, but I think it should work pretty fast overall: not too many instructions, all of them have 1 cycle of latency on modern processors. Also the stores are efficient, 2 store instructions for 20 bytes of data.

The code only uses 3 constants. If you call this function in a loop, good compilers should load all three outside of the loop and keep them in registers.

// bitwise blend according to a mask
inline void combineHigh( __m256i& vec, __m256i high, const __m256i lowMask )
{
    vec = _mm256_and_si256( vec, lowMask );
    high = _mm256_andnot_si256( lowMask, high );
    vec = _mm256_or_si256( vec, high );
}

// Store 10-bit pieces from each of the 16-bit lanes of the AVX2 vector.
// The function writes 20 bytes to the pointer.
inline void store_10x16_avx2( __m256i v, uint8_t* rdi )
{
    // Pack pairs of 10 bits into 20, into 32-bit lanes
    __m256i high = _mm256_srli_epi32( v, 16 - 10 );
    const __m256i low10 = _mm256_set1_epi32( ( 1 << 10 ) - 1 ); // Bitmask of 10 lowest bits in 32-bit lanes
    combineHigh( v, high, low10 );

    // Now the vector contains 32-bit lanes with 20 payload bits / each
    // Pack pairs of 20 bits into 40, into 64-bit lanes
    high = _mm256_srli_epi64( v, 32 - 20 );
    const __m256i low20 = _mm256_set1_epi64x( ( 1 << 20 ) - 1 ); // Bitmask of 20 lowest bits in 64-bit lanes
    combineHigh( v, high, low20 );

    // Now the vector contains 64-bit lanes with 40 payload bits / each
    // 40 bits = 5 bytes, store initial 4 bytes of the result
    _mm_storeu_si32( rdi, _mm256_castsi256_si128( v ) );

    // Shuffle the remaining 16 bytes of payload into correct positions.
    // The indices of the payload bytes are [ 0 .. 4 ] and [ 8 .. 12 ]
    // _mm256_shuffle_epi8 can only move data within 16-byte lanes
    const __m256i shuffleIndices = _mm256_setr_epi8(
        // 6 remaining payload bytes from the lower half of the vector
        4, 8, 9, 10, 11, 12,
        // 10 bytes gap, will be zeros
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        // 6 bytes gap, will be zeros
        -1, -1, -1, -1, -1, -1,
        // 10 payload bytes from the higher half of the vector
        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12
    );
    v = _mm256_shuffle_epi8( v, shuffleIndices );

    // Combine and store the final 16 bytes of payload
    const __m128i low16 = _mm256_castsi256_si128( v );
    const __m128i high16 = _mm256_extracti128_si256( v, 1 );
    const __m128i result = _mm_or_si128( low16, high16 );
    _mm_storeu_si128( ( __m128i* )( rdi + 4 ), result );
}

This code truncates unused higher 6 bits of the values.

---

If you want to saturate instead, you’ll need one more instruction, _mm256_min_epu16.

Also, if you do that, the first step of the function can use pmaddwd. Here’s the complete function which saturates the source numbers, with couple extra adjustments.

// Store 10-bit pieces from 16-bit lanes of the AVX2 vector, with saturation.
// The function writes 20 bytes to the pointer.
inline void store_10x16_avx2( __m256i v, uint8_t* rdi )
{
    const __m256i low10 = _mm256_set1_epi16( ( 1 << 10 ) - 1 );
#if 0
    // Truncate higher 6 bits; pmaddwd won't truncate, it needs zeroes in the unused higher bits.
    v = _mm256_and_si256( v, low10 );
#else
    // Saturate numbers into the range instead of truncating
    v = _mm256_min_epu16( v, low10 );
#endif

    // Pack pairs of 10 bits into 20, into 32-bit lanes
    // pmaddwd computes a[ 0 ] * b[ 0 ] + a[ 1 ] * b[ 1 ] for pairs of 16-bit lanes, making a single 32-bit number out of two pairs.
    // Initializing multiplier with pairs of [ 1, 2^10 ] to implement bit shifts + packing
    const __m256i multiplier = _mm256_set1_epi32( 1 | ( 1 << ( 10 + 16 ) ) );
    v = _mm256_madd_epi16( v, multiplier );

    // Now the vector contains 32-bit lanes with 20 payload bits / each
    // Pack pairs of 20 bits into 40 in 64-bit lanes
    __m256i low = _mm256_slli_epi32( v, 12 );
    v = _mm256_blend_epi32( v, low, 0b01010101 );
    v = _mm256_srli_epi64( v, 12 );

    // Now the vector contains 64-bit lanes with 40 payload bits / each
    // 40 bits = 5 bytes, store initial 4 bytes of the result
    _mm_storeu_si32( rdi, _mm256_castsi256_si128( v ) );

    // Shuffle the remaining 16 bytes of payload into correct positions.
    const __m256i shuffleIndices = _mm256_setr_epi8(
        // Lower half
        4, 8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        // Higher half
        -1, -1, -1, -1, -1, -1,
        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12
    );
    v = _mm256_shuffle_epi8( v, shuffleIndices );

    // Combine and store the final 16 bytes of payload
    const __m128i low16 = _mm256_castsi256_si128( v );
    const __m128i high16 = _mm256_extracti128_si256( v, 1 );
    const __m128i result = _mm_or_si128( low16, high16 );
    _mm_storeu_si128( ( __m128i* )( rdi + 4 ), result );
}

This may be slightly faster or slower overall depending on the processor, compiler, and the code calling the function, but definitely helps with code size. No one cares about binary size anymore, but CPUs have limited L1I and µop caches.

---

For completeness here’s another one that uses SSE2 and optionally SSSE3 instead of AVX2, only slightly slower in practice.

// Compute v = ( v & lowMask ) | ( high & ( ~lowMask ) ), for 256 bits of data in two registers
inline void combineHigh( __m128i& v1, __m128i& v2, __m128i h1, __m128i h2, const __m128i lowMask )
{
    v1 = _mm_and_si128( v1, lowMask );
    v2 = _mm_and_si128( v2, lowMask );
    h1 = _mm_andnot_si128( lowMask, h1 );
    h2 = _mm_andnot_si128( lowMask, h2 );
    v1 = _mm_or_si128( v1, h1 );
    v2 = _mm_or_si128( v2, h2 );
}

inline void store_10x16_sse( __m128i v1, __m128i v2, uint8_t* rdi )
{
    // Pack pairs of 10 bits into 20, in 32-bit lanes
    __m128i h1 = _mm_srli_epi32( v1, 16 - 10 );
    __m128i h2 = _mm_srli_epi32( v2, 16 - 10 );
    const __m128i low10 = _mm_set1_epi32( ( 1 << 10 ) - 1 );
    combineHigh( v1, v2, h1, h2, low10 );

    // Pack pairs of 20 bits into 40, in 64-bit lanes
    h1 = _mm_srli_epi64( v1, 32 - 20 );
    h2 = _mm_srli_epi64( v2, 32 - 20 );
    const __m128i low20 = _mm_set1_epi64x( ( 1 << 20 ) - 1 );
    combineHigh( v1, v2, h1, h2, low20 );

#if 1
    // 40 bits is 5 bytes, for the final shuffle we use pshufb instruction from SSSE3 set
    // If you don't have SSSE3, below under `#else` there's SSE2-only workaround.
    const __m128i shuffleIndices = _mm_setr_epi8(
        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1 );
    v1 = _mm_shuffle_epi8( v1, shuffleIndices );
    v2 = _mm_shuffle_epi8( v2, shuffleIndices );
#else
    // SSE2-only version of the above, uses 8 instructions + 2 constants to emulate 2 instructions + 1 constant
    // Need two constants because after this step we want zeros in the unused higher 6 bytes.
    h1 = _mm_srli_si128( v1, 3 );
    h2 = _mm_srli_si128( v2, 3 );
    const __m128i low40 = _mm_setr_epi8( -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 );
    const __m128i high40 = _mm_setr_epi8( 0, 0, 0, 0, 0, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0 );
    const __m128i l1 = _mm_and_si128( v1, low40 );
    const __m128i l2 = _mm_and_si128( v2, low40 );
    h1 = _mm_and_si128( h1, high40 );
    h2 = _mm_and_si128( h2, high40 );
    v1 = _mm_or_si128( h1, l1 );
    v2 = _mm_or_si128( h2, l2 );
#endif

    // Now v1 and v2 vectors contain densely packed 10 bytes / each.
    // Produce final result: 16 bytes in the low part, 4 bytes in the high part
    __m128i low16 = _mm_or_si128( v1, _mm_slli_si128( v2, 10 ) );
    __m128i high16 = _mm_srli_si128( v2, 6 );
    // Store these 20 bytes with 2 instructions
    _mm_storeu_si128( ( __m128i* )rdi, low16 );
    _mm_storeu_si32( rdi + 16, high16 );
}

Answer 2

In a loop, you may want to use partially-overlapping stores that write past the end of the 20 byte destination for each vector of source data. That saves the work of shuffling data across the 16-byte boundary to set up for 16 + 4 byte stores.

(@Soont's updated answer with one vmovd and one vmovdqu store is very good and only has 2 total shuffle uops including vpshufb and vextracti128. When I initially wrote this, we hadn't yet thought of a good way to avoid storing outside the 20 bytes without spending more shuffle uops which would create a bottleneck worse than the front-end. But vmovdqu + vextracti128 mem, ymm, 1 (2 uops not micro-fused) is still slightly cheaper: 3 uops after the vpshufb instead of 4.)

Or unrolling could be good for large arrays, LCM(20,16) = 80, so with a large unroll (and different shuffle-control vectors for each position within that) you could be doing only aligned 16-byte stores. But that might take a lot of shuffling, including between source chunks probably with palignr.

---

Example of two overlapping 16-byte stores

Use this as a loop body where overwriting past 20 bytes is ok.

#include <immintrin.h>
#include <stdint.h>

// Store 10-bit pieces from each of the 16-bit lanes of the AVX2 vector.
// The function writes 20 useful bytes to the pointer
// but actually steps on data out to 26 bytes from dst
void pack10bit_avx2_store26( __m256i v, uint8_t* dst)
{
    // clear high garbage if elements aren't already zero-extended   
    //v = _mm256_and_si256(v, _mm256_set1_epi16( (1<<10)-1) );

    ... prep data somehow; pmaddwd + a couple shifts is good for throughput

    // Now the vector contains 64-bit lanes with 40 payload bits / each; 40 bits = 5 bytes.
    // Shuffle these bytes into a very special order.
    // Note _mm256_shuffle_epi8 can only move data within 16-byte lanes.
    const __m256i shuffleIndices = _mm256_setr_epi8(
        // 6 bytes gap with zeros
        // Pack the two 5-byte chunks into the bottom of each 16-byte lane
        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1,

        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1);
    v = _mm256_shuffle_epi8(v, shuffleIndices );

    // Split the vector into halves
    __m128i low16 = _mm256_castsi256_si128( v );
    _mm_storeu_si128( ( __m128i* )dst, low16 );        // vmovdqu      mem, xmm

    __m128i high16 = _mm256_extracti128_si256( v, 1 );
    _mm_storeu_si128( ( __m128i* )(dst+10), high16 );   // vextracti128 mem, ymm, 1

    // An AVX-512 masked store could avoid writing past the end
}

We can see how it might inline into a loop by compiling it to a stand-alone function (https://godbolt.org/z/8T7KhT).

# clang -O3 -march=skylake
pack10bit_avx2(long long __vector(4), unsigned char*):
       # vpand  commented out
        vpmaddwd        ymm0, ymm0, ymmword ptr [rip + .LCPI0_0]
         ... # work in progress, original PMADDWD idea ignored some limitations!  See Soonts' answer

        vpshufb ymm0, ymm0, ymmword ptr [rip + .LCPI0_1] # ymm0 = ymm0[0,1,2,3,4,8,9,10,11,12],zero,zero,zero,zero,zero,zero,ymm0[16,17,18,19,20,24,25,26,27,28],zero,zero,zero,zero,zero,zero
        vmovdqu xmmword ptr [rdi], xmm0
        vextracti128    xmmword ptr [rdi + 10], ymm0, 1

        vzeroupper               # overhead that goes away when inlining into a loop
        ret

In a loop, compilers would load those 2 vector constants into registers, hopefully using broadcast-loads.

Unlike some wider integer multiplies or horizontal add, vpmaddwd is handed efficiently, as a single uop with 5 cycle latency. https://uops.info/

The vextracti128 store can't micro-fuse on Intel, but unlike vpextrd there's no shuffle uop involved. Just store-address and store-data. Zen2 also runs it as 2 uops, with throughput of one per 2 cycles unfortunately. (Worse than Zen1).

Before Ice Lake, both Intel and AMD can run 1 store per clock.

---

If you do actually want the packed data back in registers, you might want @Soont's original shuffles using palignr, or you could do a block of this and then some reloads. Latency would be higher (especially because of store-forwarding stalls on the reloads), but if your block is several registers worth of data then that should overlap or even hide the latency, maybe giving the stores time to even commit to L1d and not cause a stall when reloaded.

---

BMI2 pext

uint64_t packed = _pext_u64(x, 0x03FF03FF03FF03FF);

Maybe good for scalar cleanup or a short chunk of 4 pixels or whatever. This leaves you with the problem of doing a 5-byte store (or 8-byte store with trailing zeros). Beware of strict-aliasing and alignment if using this, e.g. use memcpy to get unaligned may-alias data into a uint64_t, or make an __attribute__((aligned(1),may_alias)) typedef.

pext is very efficient on Intel (1 uop, 3c latency), but very bad on AMD, much worse than just using the low part of one SIMD step.

---

AVX-512

AVX512VBMI (Ice Lake) would give you vpermb (lane crossing) instead of vpshufb. (AVX512BW for vpermw on Skylake-X / Cascade Lake would require you to have already combined into an even number of bytes, and it's 2 uops even on Ice Lake where vpermb is 1, so it's pretty bad.) vpermb could set up for a single unaligned 32-byte store (with 20 useful bytes), which you overlap in a loop.

AVX-512 stores can be efficiently masked to not actually overwrite past the end, e.g. using dword masking. vmovdqu32 [rdi]{k}, ymm0 is 1 uop on Skylake-X. But AVX2 vmaskmovd is a few uops even on Intel, and extremely expensive on AMD, so you don't want to do that. And dword masking only works if you have all 20 bytes ready for one store, otherwise you need at least 16-bit granularity.

Other AVX-512 instructions: VBMI vpmultishiftqb, a parallel bitfield extract, seems like it might be useful, but it can only write aligned 8-bit destination chunks from unaligned but contiguous source chunks. I don't think that's better than what we can do with variable-shifts and rotates. vpmultishiftqb would let us unpack this format (inverse of this function) in probably 2 instructions: 1 shuffle (such as vpexpandb or vpermb) to put the needed data into each qword in the vector, and one multishift to grab the right 10-bit field for the bottom of each word.

AVX-512 has variable-count shifts and rotates, including with word (16-bit) granularity, so that would be an option instead of vpmaddwd for the first step. Using shifts ignores high garbage for free. It has lower latency, and merge-masking for the immediate version can replace the need for a control vector. (But then you need a mask constant).

With masking the latency is 3 cycles, vs 1 without, and AVX-512 makes it about as efficient to broadcast a control vector from an immediate as to mov reg,imm / kmov kreg, reg. e.g. mov reg,imm / vpbroadcastd ymm, reg (1 uop). Merge-masking also constrains the optimizer to overwrite the destination register instead of copy-and-shift, although that shouldn't matter here if the optimizer is smart. Neither way lets the load of the data fold into a memory source operand for the shift: sllvw can only take the counts from memory, and sllw needs to merge into the original in a register.

Shifts can run on ports 0 or 1 on Intel (and AMD doesn't support AVX-512). Or only port 0 for 512-bit uops, shutting down port 1 for any vector-ALU uop while any 512-bit uops are in flight. So there's a potential throughput bottleneck on port 0 for a __m512i version of this, but for 256-bit there are enough other uops (shuffle and store, and presumably loop overhead if doing this for an array of data) that this should be fairly evenly distributed.

This shift part (before _mm256_permutexvar_epi8) only requires AVX-512BW (+VL), and will work on Skylake-X. It leaves the data in the same place as other methods, so is a drop-in replacement you can mix and match with various strategies.

// Ice Lake.  Could work on __m512i but then shifts could only run on p0, not p0/p1,
  //  and almost every store would be a cache line split.
inline void store_10x16_avx512vbmi( __m256i v, uint8_t* dst )
{
// no _mm256_and_si256 needed, we safely ignore high bits
   // v = [ ?(6) ... B[9:0] | ?(6) ... A[9:0] ] repeated
   v = _mm256_sllv_epi16(v, _mm256_set1_epi32((0<<16) | 6));  // alternative: simple repeated-pattern control vector
      // v =  _mm256_mask_slli_epi16(v, 0x5555, v, 6);   // merge-masking, updating only elements 0,2, etc.
   // v = [ ?(6) ... B[9:0] | A[9:0] ... 0(6) ] repeated
   v = _mm256_rolv_epi32(v, _mm256_set1_epi64x(((32ULL-6)<<32) | 6));  // top half right, bottom half left
   // v = [ 0(6) .. ?(6) .. D[9:0] | C[9:0] | B[9:0] | A[9:0] ... 0(12) ] repeated
   v = _mm256_srli_epi64(v, 12);    // 40 bit chunks at the bottom of each qword

   const __m256i permb = _mm256_setr_epi8( 0, 1, 2, 3, 4,   8, 9,10,11,12,
                                          16,17,18,19,20,  24,25,26,27,28,
                                          28,28,28,28,28,28,28,28,28,28,28,28 );
    // repeat last byte as filler.  vpermb can't zero (except by maskz) but we can do a masked store
   v = _mm256_permutexvar_epi8(v, permb);  // AVX512_VBMI
   _mm256_mask_storeu_epi32( dst, 0x1F, v);  // 32-bit masking granularity in case that's cheaper for HW.  20 bytes = 5 dwords.
}

Compiles like so (Godbolt):

# clang -O3 -march=icelake-client.  GCC is essentially the same.
store_10x16_avx512vbmi(long long __vector(4), unsigned char*):
        vpsllvw ymm0, ymm0, ymmword ptr [rip + .LCPI0_0]
        vprolvd ymm0, ymm0, ymmword ptr [rip + .LCPI0_1]
        vpsrlq  ymm0, ymm0, 12
        vpermb  ymm0, ymm0, ymmword ptr [rip + .LCPI0_2]
        mov     al, 31           # what the heck, clang? partial register false dependency for no reason!
        kmovd   k1, eax
        vmovdqu32       ymmword ptr [rdi] {k1}, ymm0
      # vzeroupper not needed because the caller was using __m256i args.  GCC omits it.
        ret

Even if you use the same shift constant vector twice to make the compiler keep it around in a register (instead of use directly from a memory source operand), it still chooses to load it from memory instead of mov eax,6 / vpbroadcast ymm1, eax or something. This saves 1 uop at the cost of needing the constant in .rodata. To be fair, we do need other constants probably in the same cache line, but the way GCC wastes space they don't all fit in one cache line! clang notices the pattern and uses a vpbroadcastd or q load, gcc wastefully loads a full 32 bytes. (kmov k1, [mem] is 3 front-end uops so it wouldn't save a uop to load mask constants from memory.)

Using _mm256_mask_slli_epi16(v, 0x5555, v, 6), clang optimizes it back into vpsllvw ymm0, ymm0, ymmword ptr [rip + .LCPI0_0] with the same 6,0 repeating constant. So I guess that's a good sign I got it right. But GCC compiles as written:

store_10x16_avx512vbmi(long long __vector(4), unsigned char*):
        mov     eax, 21845
        kmovw   k1, eax
        vpsllw  ymm0{k1}, ymm0, 6
        vprolvd ymm0, ymm0, YMMWORD PTR .LC0[rip]
        mov     eax, 31
        kmovb   k2, eax
        vpsrlq  ymm0, ymm0, 12
        vpermb  ymm0, ymm0, YMMWORD PTR .LC1[rip]
        vmovdqu32       YMMWORD PTR [rdi]{k2}, ymm0
        ret

---

_mm256_sllv_epi16 requires AVX-512BW and AVX-512VL. rolv_epi32 only requires AVX-512VL. (Or just AVX-512F for the 512-bit version.) Rotates only come in 32 and 64 element sizes, not 16, but AVX-512 does extend variable-shift granularity down to 16 (from 32 or 64 in AVX2).

vpcompressb [rdi]{k1}, ymm0 (AVX512VBMI = Ice Lake and later) would be an alternative to vpermb + store to pack bytes at the bottom of a register (like BMI2 pext but for vector elements instead of bits in a scalar register). But it's actually more expensive: 6 uops on Ice Lake, with one per 6c throughput. (vpcompressd is not as bad).

Even vpcompressb into a vector register is 2 uops, so for a constant shuffle-control it's better to load a vector constant for vpermb, unless cache misses for control vectors is a problem, e.g. if you're only doing this once every so often then let the HW process a k mask instead of a load.

---

AVX-512 without VBMI: 2x 16-byte stores without exceeding the 20-byte range

   ...  // same setup as usual, leaving 40-bit chunks at the bottom of each qword
 
    const __m256i shuffleIndices = _mm256_setr_epi8(
        // 6 bytes gap with zeros
        // Pack the two 5-byte chunks into the bottom of each 16-byte lane
        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1,

        0, 1, 2, 3, 4,
        8, 9, 10, 11, 12,
        -1, -1, -1, -1, -1, -1);
    v = _mm256_shuffle_epi8(v, shuffleIndices );

    // Split the vector into halves
    __m128i low16 = _mm256_castsi256_si128( v );
    _mm_storeu_si128( ( __m128i* )dst, low16 );        // vmovdqu      mem, xmm  no masking

    // An AVX-512BW masked store avoiding writing past the end costs more instructions (and back-end uops), same front-end uops
    __m128i high16 = _mm256_extracti128_si256( v, 1 );  // vextracti128 xmm, ymm, 1
    _mm_mask_storeu_epi8( dst+10, 0x3FF, high16 );      // vmovdqu8 [mem]{k}, xmm

This needs vextracti128 xmm, ymm, 1 to set up for vmovdqu8. Unlike with writing 26 bytes, we can't extract directly to memory. There is no vextracti8x16, only vextracti32x4 and 64x2 (and 32x8 / 64x4 256-bit extracts). We need byte-granularity masking but can't get it with an instruction that extracts directly to memory, only via a shuffle (vextract into a register) and then vmovdqu8.

So the asm we get is

# clang
...        vpshufb result in YMM0
        vmovdqu      [rdi], xmm0             # same as before
        vextracti128    xmm0, ymm0, 1        # 1 shuffle uop
        mov     ax, 1023
        kmovd   k1, eax                         # will be hoisted
        vmovdqu8     [rdi + 10] {k1}, xmm0   # 1 micro-fused uop

Since vextracti128 [mem], ymm, 1 was 2 front-end uops anyway, this doesn't hurt front-end throughput. (It does create more pressure on back-end execution ports, thanks to the shuffle uop).