Function byte_dists_vertical

cpp/src/quantize/multi_codebook.hpp:765–806

block size, number of dimensions of each vector

template<int B, int D> // block size, number of dimensions of each vector
inline void byte_dists_vertical(const uint8_t* X, const uint8_t* q,
    uint16_t* dists_out, int64_t N)
{
    static constexpr int packet_width = 32; // how many vecs we operate on at once
    static constexpr int nstripes = B / packet_width; // # of rows of 32B per block
    static_assert(B % packet_width == 0, "B must be a multiple of packet_width");
    static_assert(B > 0, "B must be > 0");
    const int64_t nblocks = N / B * 2;
    assert(N % B == 0);

    // we assume that pairs of bytes are from the same vector for our
    // maddubs; so pretending q points to int16s makes broadcasting pairs work
    const uint16_t* q16 = reinterpret_cast<const uint16_t*>(q);

    __m256i accumulators[nstripes];

    for (int64_t b = 0; b < nblocks; b++) { // for each block
        for (int i = 0; i < nstripes; i++) {
            accumulators[i] = _mm256_setzero_si256(); // zero dists
        }
        for (int j = 0; j < D / 2; j++) { // for each pair of dimensions
            auto q_broadcast = _mm256_set1_epi16(q16[j]);
            for (int i = 0; i < nstripes; i++) { // for each stripe
                // auto x_col = _mm256_load_si256((__m256i*)X);
                auto x_col = stream_load_si256i((__m256i*)X);
                X += packet_width;

                auto diffs = _mm256_sub_epi8(q_broadcast, x_col);
                auto abs_diffs = _mm256_abs_epi8(diffs);
                auto prods = _mm256_maddubs_epi16(abs_diffs, abs_diffs);
                accumulators[i] = _mm256_adds_epi16(accumulators[i], prods);
            }
        }
        for (uint8_t i = 0; i < nstripes; i++) { // for each stripe
            _mm256_store_si256((__m256i*)dists_out, accumulators[i]);
            // we compute 16 dists because we assume that adjacent pairs of
            // bytes belong to the same vector; this is necessary for the
            // maddubs to be meaningful
            dists_out += packet_width / 2;
        }
    }
}

// just sums up the inputs and stores results; upper bound on how fast
// any of these functions could be