AVX2 implementation

The following C++ template generates an implementation. The most generic one would scan all mask bits (MAX_MASK_BITS=32), but when we know the maximum number of bits set in masks, we may limit the number of iterations accordingly.

The parameter EARLY_EXIT decides whether to check in every iteration if all mask bits were scanned on not. It adds runtime overhead, and in the case we don't know how diverse masks are, this check won't help. If we know that the number of bits in different masks are similar, this check may be a win.

template <int MAX_MASK_BITS, bool EARLY_EXIT>
void avx2_pext_u32_reference(const uint32_t* data_arr, const uint32_t* mask_arr, uint32_t* out_arr, size_t n) {
    static_assert(MAX_MASK_BITS > 0);
    static_assert(MAX_MASK_BITS <= 32);

    const __m256i one  = _mm256_set1_epi32(1);
    const __m256i zero = _mm256_set1_epi32(0);
    for (size_t i=0; i < n; i += 8) {
        __m256i data = _mm256_loadu_si256((const __m256i*)(&data_arr[i]));
        __m256i mask = _mm256_loadu_si256((const __m256i*)(&mask_arr[i]));
        __m256i out  = _mm256_set1_epi32(0);

        __m256i bit = one;

        for (int j=0; j < MAX_MASK_BITS; j++) {
            // 1. isolate the first non-zero bit set of mask
            const __m256i m0 = _mm256_sub_epi32(mask, one);
            const __m256i m1 = _mm256_and_si256(mask, m0);
            const __m256i m2 = _mm256_xor_si256(mask, m1);

            // 2. isolate that bit from data word
            const __m256i d0 = _mm256_and_si256(data, m2);

            // 3. move that bit on the next position in out
            const __m256i o0 = _mm256_min_epu32(d0, bit);
            out = _mm256_or_si256(out, o0);

            // 4. reset selected bit in mask (we already done it)
            mask = m1;

            // 5. advance to the next bit in out
            bit = _mm256_add_epi32(bit, bit);

            // 6. all are zeros?
            if (EARLY_EXIT && _mm256_testc_si256(zero, mask)) {
                break;
            }
        }

        _mm256_storeu_si256((__m256i*)(&out_arr[i]), out);
    }
}

AVX-512 implementation

Solution for AVX-512 can use the above trick, but it is simpler to directly use mask registers. There is instruction VPTESTMD that sets a register mask based on bit-and of two operands. This allows us to produce element-wise masks from m-th bits.

Then conditional VPORD can be used to update the resulting word.

The following fragment from AVX-512 implementation shows these steps.

// 1. isolate the first bit set of mask
//                                                 mask = [0101_1001_1100_0000|...]
const __m512i m0 = _mm512_sub_epi32(mask, one); // m0   = [0101_1001_1011_1111|...]
const __m512i m1 = _mm512_and_si512(mask, m0);  // m1   = [0101_1001_1000_0000|...]
const __m512i m2 = _mm512_xor_si512(mask, m1);  // m2   = [0000_0000_0100_0000|...]

// 2. get m-th data bits in a mask register
const __mmask64 m = _mm512_test_epi32_mask(data, m2);

// 3. set k-th bit when m-th bit is set
out = _mm512_mask_or_epi32(out, m, out, bit);

// 4. reset selected bit in mask
mask = m1;

// 5. the next bit to set
bit = _mm512_add_epi32(bit, bit);

AVX2 implementation

While the shift-less method is feasible, we found an easier way to update the result word. We use VPCMPEQD that fills a 32-bit word with zeros or ones, depending on comparison result. We isolate the k-th bit of data and compare it with zero.

This yield a negated mask, which we use to filter out m-th bits.

The following snippet from AVX2 implementation shows the order of operations.

// 1. isolate the first non-zoro bit set of mask (at m)
//                                                 mask = [0101_1001_1100_0000|0000_1110_1100_1000|...]
const __m256i m0 = _mm256_sub_epi32(mask, one); // m0   = [0101_1001_1011_1111|0000_1110_1100_0111|...]
const __m256i m1 = _mm256_and_si256(mask, m0);  // m1   = [0101_1001_1000_0000|0000_1110_1100_0000|...]
const __m256i m2 = _mm256_xor_si256(mask, m1);  // m2   = [0000_0000_0100_0000|0000_0000_0000_1000|...]

// 2. isolate k-th bit from data                   data = [1100_0000_1111_1110|0000_0000_1000_0000|...]
//                                                  bit = [0000_0000_0001_0000]0000_0000_0001_0000|...]
const __m256i d0 = _mm256_and_si256(data, bit); //   d0 = [0000_0000_0001_0000|0000_0000_0000_0000|...]
                                                //                      ^                   ^
// 4. fill word with *negation* of data bit
const __m256i d1 = _mm256_cmpeq_epi32(d0, zero);//   d1 = [0000_0000_0000_0000|1111_1111_1111_1111|...]

// 5. keep the mask bit, iff data[k] == 1
const __m256i m3 = _mm256_andnot_si256(d1, m2); //   m3 = [0000_0000_0100_0000|0000_0000_0000_0000|...]

// 6. update the out
out = _mm256_or_si256(out, m3);
mask = m1;

// 7. the next bit in data to check
bit = _mm256_add_epi32(bit, bit);

AVX-512 implementation

AVX-512 code uses the same approach as PEXT implementation. We transfer k-th bits from data word to a mask register and then use masked VPTERNLOGD to conditionally bit-or the m-th bits into the result words.

// 1. get k-th data bits in a mask register
const __mmask64 m = _mm512_test_epi32_mask(data, bit);

// 2. isolate the first non-zoro bit set of mask (at m)

//                                                 mask = [0101_1001_1100_0000|0000_1110_1100_1000|...]
const __m512i m0 = _mm512_sub_epi32(mask, one); // m0   = [0101_1001_1011_1111|0000_1110_1100_0111|...]
const __m512i m1 = _mm512_and_si512(mask, m0);  // m1   = [0101_1001_1000_0000|0000_1110_1100_0000|...]

// 3. out = out | (mask ^ m1)
out = _mm512_mask_ternarylogic_epi32(out, m, mask, m1, 0xf6);

// 4. the next mask value
mask = m1;

// 5. the next bit in data to update
bit = _mm512_add_epi32(bit, bit);