compares quantization of means vs 1-2 dim per bit quantize

1vs2bitvs4_quantize_test.cpp

// bucket_stats.cpp
#include <array>
#include <vector>
#include <unordered_map>
#include <random>
#include <iostream>
#include <algorithm>
#include <chrono>
#include <cstdint>
#include <numeric>
#include <cassert>
#include <bitset>
#include <ctime>


using namespace std;

constexpr auto quantize_byte_1dim(const array<float, 64>& v, size_t start, float threshold) {
    return (v[start] > threshold ? 1 : 0) |
           (v[start + 1] > threshold ? 1 << 1 : 0) |
           (v[start + 2] > threshold ? 1 << 2 : 0) |
           (v[start + 3] > threshold ? 1 << 3 : 0) |
           (v[start + 4] > threshold ? 1 << 4 : 0) |
           (v[start + 5] > threshold ? 1 << 5 : 0) |
           (v[start + 6] > threshold ? 1 << 6 : 0) |
           (v[start + 7] > threshold ? 1 << 7 : 0);
}

constexpr bool mean_of_2dim_thresh(float f1, float f2, float threshold) {
    return ((f1 + f2) / 2) > threshold;
}

constexpr bool mean_of_4dim_thresh(float f1, float f2, float f3, float f4, float threshold) {
    return ((f1 + f2 + f3 + f4) / 4) > threshold;
}

constexpr auto quantize_byte_2dim(const array<float, 64>& v, size_t start, float threshold) {
    return (mean_of_2dim_thresh(v[start], v[start+1], threshold) ? 1 : 0) |
           (mean_of_2dim_thresh(v[start+2], v[start+3], threshold) ? 1 << 1 : 0) |
           (mean_of_2dim_thresh(v[start+4], v[start+5], threshold) ? 1 << 2 : 0) |
           (mean_of_2dim_thresh(v[start+6], v[start+7], threshold) ? 1 << 3 : 0) |
           (mean_of_2dim_thresh(v[start+8], v[start+9], threshold) ? 1 << 4 : 0) |
           (mean_of_2dim_thresh(v[start+10], v[start+11], threshold) ? 1 << 5 : 0) |
           (mean_of_2dim_thresh(v[start+12], v[start+13], threshold) ? 1 << 6 : 0) |
           (mean_of_2dim_thresh(v[start+14], v[start+15], threshold) ? 1 << 7 : 0);
}

constexpr auto quantize_byte_4dim(const array<float, 64>& v, size_t start, float threshold) {
    return (mean_of_4dim_thresh(v[start], v[start+1], v[start+2], v[start+3], threshold) ? 1 : 0) |
           (mean_of_4dim_thresh(v[start+4], v[start+5], v[start+6], v[start+7], threshold) ? 1 << 1 : 0) |
           (mean_of_4dim_thresh(v[start+8], v[start+9], v[start+10], v[start+11], threshold) ? 1 << 2 : 0) |
           (mean_of_4dim_thresh(v[start+12], v[start+13], v[start+14], v[start+15], threshold) ? 1 << 3 : 0) |
           (mean_of_4dim_thresh(v[start+16], v[start+17], v[start+18], v[start+19], threshold) ? 1 << 4 : 0) |
           (mean_of_4dim_thresh(v[start+20], v[start+21], v[start+22], v[start+23], threshold) ? 1 << 5 : 0) |
           (mean_of_4dim_thresh(v[start+24], v[start+25], v[start+26], v[start+27], threshold) ? 1 << 6 : 0) |
           (mean_of_4dim_thresh(v[start+28], v[start+29], v[start+30], v[start+31], threshold) ? 1 << 7 : 0);
}


constexpr auto quantize_vector_64_1b(const array<float, 64>& v, float threshold = 0.5f) {

    array<uint8_t, 8> q{};

    q[0] = quantize_byte_1dim(v, 0, threshold);

    q[1] = quantize_byte_1dim(v, 8, threshold);

    q[2] = quantize_byte_1dim(v, 16, threshold);
    q[3] = quantize_byte_1dim(v, 24, threshold);
    q[4] = quantize_byte_1dim(v, 32, threshold);
    q[5] = quantize_byte_1dim(v, 40, threshold);
    q[6] = quantize_byte_1dim(v, 48, threshold);
    q[7] = quantize_byte_1dim(v, 56, threshold);

    return q;
}

constexpr auto quantize_vector_64_2b(const array<float, 64>& v, float threshold = 0.5f) {

    array<uint8_t, 4> q{};

    q[0] = quantize_byte_2dim(v, 0, threshold);

    q[1] = quantize_byte_2dim(v, 16, threshold);
    q[2] = quantize_byte_2dim(v, 32, threshold);
    q[3] = quantize_byte_2dim(v, 48, threshold);

    return q;
}

constexpr auto quantize_vector_64_4b(const array<float, 64>& v, float threshold = 0.5f) {
    array<uint8_t, 2> q{};
    q[0] = quantize_byte_4dim(v, 0, threshold);
    q[1] = quantize_byte_4dim(v, 32, threshold);
    return q;
}

constexpr uint64_t byte_arr_to_64bit(const std::array<uint8_t, 8>& q) {
    return (uint64_t(q[0]) << 0)  |
           (uint64_t(q[1]) << 8)  |
           (uint64_t(q[2]) << 16) |
           (uint64_t(q[3]) << 24) |
           (uint64_t(q[4]) << 32) |
           (uint64_t(q[5]) << 40) |
           (uint64_t(q[6]) << 48) |
           (uint64_t(q[7]) << 56);
}

constexpr uint32_t byte_arr_to_32bit(const std::array<uint8_t, 4>& q) {
    return (uint32_t(q[0]) << 0)  |
           (uint32_t(q[1]) << 8)  |
           (uint32_t(q[2]) << 16) |
           (uint32_t(q[3]) << 24);
}

constexpr uint16_t byte_arr_to_16bit(const std::array<uint8_t, 2>& q) {
    return (uint16_t(q[0]) << 0)  |  (uint16_t(q[1]) << 8);
}

struct Vec {
    uint64_t innerVec = 0;
    std::array<float, 64> fVec;
};


int main(int argc, char const *argv[])
{
    std::mt19937_64 rng(time(nullptr));
    std::normal_distribution<float> ndist(0.0f, 1.0f);

    std::array<float, 64> vec;

    for (size_t i = 0; i < 10; ++i) {
        // generate one vector
        for (size_t d = 0; d < 64; ++d) vec[d] = ndist(rng);
        const auto generated = quantize_vector_64_1b(vec);
        std::cout << "-------\n";
        std::cout << byte_arr_to_64bit(generated) << "\n";
    }

    for (size_t i = 0; i < 10; ++i) {
        // generate one vector
        for (size_t d = 0; d < 64; ++d) vec[d] = ndist(rng);
        const auto generated = quantize_vector_64_2b(vec);
        std::cout << "-------\n";
        std::cout << byte_arr_to_32bit(generated) << "\n";
    }

    {
        std::cout << "Starting map test for collision 64 bit test\n";
        std::unordered_map<uint64_t, size_t> map64coll;

        for (size_t i = 0; i < 10000000; ++i)
        {
            for (size_t d = 0; d < 64; ++d) vec[d] = ndist(rng);
            const auto generated = quantize_vector_64_1b(vec);
            const auto mapKey = byte_arr_to_64bit(generated);
            auto foundIt = map64coll.find(mapKey);
            if (foundIt == map64coll.end()) {
                map64coll[mapKey] = 1;
            } else {
                map64coll[mapKey] += 1;
            }
        }

        size_t total_collided_buckets = 0;
        for (const auto& p : map64coll) {
            if (p.second > 1) {
                std::cout << "key=" << p.first << ", count=" << p.second << "\n";
                ++total_collided_buckets;
            }
        }
        std::cout << "Total 64bit shared buckets=" << total_collided_buckets << "\n";
    }

    {
        std::cout << "Starting map test for collision 32 bit test\n";
        std::unordered_map<uint32_t, size_t> map32coll;

        for (size_t i = 0; i < 10000000; ++i)
        {
            for (size_t d = 0; d < 64; ++d) vec[d] = ndist(rng);
            const auto generated = quantize_vector_64_2b(vec);
            const auto mapKey = byte_arr_to_32bit(generated);
            auto foundIt = map32coll.find(mapKey);
            if (foundIt == map32coll.end()) {
                map32coll[mapKey] = 1;
            } else {
                map32coll[mapKey] += 1;
            }
        }

        size_t total_collided_buckets = 0;
        for (const auto& p : map32coll) {
            if (p.second > 1) {
                std::cout << "key=" << p.first << ", count=" << p.second << "\n";
                ++total_collided_buckets;
            }
        }
        std::cout << "Total 32bit shared buckets=" << total_collided_buckets << "\n";
    }

    {
        std::cout << "Starting map test for collision 16 bit test\n";
        std::unordered_map<uint16_t, size_t> map16coll;

        for (size_t i = 0; i < 10000000; ++i)
        {
            for (size_t d = 0; d < 64; ++d) vec[d] = ndist(rng);
            const auto generated = quantize_vector_64_4b(vec);
            const auto mapKey = byte_arr_to_16bit(generated);
            auto foundIt = map16coll.find(mapKey);
            if (foundIt == map16coll.end()) {
                map16coll[mapKey] = 1;
            } else {
                map16coll[mapKey] += 1;
            }
        }

        size_t total_collided_buckets = 0;
        for (const auto& p : map16coll) {
            if (p.second > 1) {
                std::cout << "key=" << p.first << ", count=" << p.second << "\n";
                ++total_collided_buckets;
            }
        }
        std::cout << "Total 16bit shared buckets=" << total_collided_buckets << "\n";
    }
    return 0;

/***
 * key=2193623298, count=2
key=708847752, count=2
key=3288342529, count=3
key=1378353681, count=2
key=1109415936, count=2
key=402917448, count=2
key=71573584, count=3
key=677388800, count=3
key=20972563, count=2
key=84018210, count=3
key=84017156, count=17
key=1077946436, count=2
key=2550137024, count=5
Total 32bit shared buckets=1076932
 * */
}