hnswlib 0.1.0

data/ext/hnswlib/src/hnswlib.h

@@ -0,0 +1,108 @@
+#pragma once
+#ifndef NO_MANUAL_VECTORIZATION
+#ifdef __SSE__
+#define USE_SSE
+#ifdef __AVX__
+#define USE_AVX
+#endif
+#endif
+#endif
+
+#if defined(USE_AVX) || defined(USE_SSE)
+#ifdef _MSC_VER
+#include <intrin.h>
+#include <stdexcept>
+#else
+#include <x86intrin.h>
+#endif
+
+#if defined(__GNUC__)
+#define PORTABLE_ALIGN32 __attribute__((aligned(32)))
+#else
+#define PORTABLE_ALIGN32 __declspec(align(32))
+#endif
+#endif
+
+#include <queue>
+#include <vector>
+#include <iostream>
+#include <string.h>
+
+namespace hnswlib {
+    typedef size_t labeltype;
+
+    template <typename T>
+    class pairGreater {
+    public:
+        bool operator()(const T& p1, const T& p2) {
+            return p1.first > p2.first;
+        }
+    };
+
+    template<typename T>
+    static void writeBinaryPOD(std::ostream &out, const T &podRef) {
+        out.write((char *) &podRef, sizeof(T));
+    }
+
+    template<typename T>
+    static void readBinaryPOD(std::istream &in, T &podRef) {
+        in.read((char *) &podRef, sizeof(T));
+    }
+
+    template<typename MTYPE>
+    using DISTFUNC = MTYPE(*)(const void *, const void *, const void *);
+
+
+    template<typename MTYPE>
+    class SpaceInterface {
+    public:
+        //virtual void search(void *);
+        virtual size_t get_data_size() = 0;
+
+        virtual DISTFUNC<MTYPE> get_dist_func() = 0;
+
+        virtual void *get_dist_func_param() = 0;
+
+        virtual ~SpaceInterface() {}
+    };
+
+    template<typename dist_t>
+    class AlgorithmInterface {
+    public:
+        virtual void addPoint(const void *datapoint, labeltype label)=0;
+        virtual std::priority_queue<std::pair<dist_t, labeltype >> searchKnn(const void *, size_t) const = 0;
+
+        // Return k nearest neighbor in the order of closer fist
+        virtual std::vector<std::pair<dist_t, labeltype>>
+            searchKnnCloserFirst(const void* query_data, size_t k) const;
+
+        virtual void saveIndex(const std::string &location)=0;
+        virtual ~AlgorithmInterface(){
+        }
+    };
+
+    template<typename dist_t>
+    std::vector<std::pair<dist_t, labeltype>>
+    AlgorithmInterface<dist_t>::searchKnnCloserFirst(const void* query_data, size_t k) const {
+        std::vector<std::pair<dist_t, labeltype>> result;
+
+        // here searchKnn returns the result in the order of further first
+        auto ret = searchKnn(query_data, k);
+        {
+            size_t sz = ret.size();
+            result.resize(sz);
+            while (!ret.empty()) {
+                result[--sz] = ret.top();
+                ret.pop();
+            }
+        }
+
+        return result;
+    }
+
+}
+
+#include "space_l2.h"
+#include "space_ip.h"
+#include "bruteforce.h"
+#include "hnswalg.h"

data/ext/hnswlib/src/space_ip.h

@@ -0,0 +1,282 @@
+#pragma once
+#include "hnswlib.h"
+
+namespace hnswlib {
+
+    static float
+    InnerProduct(const void *pVect1, const void *pVect2, const void *qty_ptr) {
+        size_t qty = *((size_t *) qty_ptr);
+        float res = 0;
+        for (unsigned i = 0; i < qty; i++) {
+            res += ((float *) pVect1)[i] * ((float *) pVect2)[i];
+        }
+        return (1.0f - res);
+
+    }
+
+#if defined(USE_AVX)
+
+// Favor using AVX if available.
+    static float
+    InnerProductSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float PORTABLE_ALIGN32 TmpRes[8];
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+        size_t qty16 = qty / 16;
+        size_t qty4 = qty / 4;
+
+        const float *pEnd1 = pVect1 + 16 * qty16;
+        const float *pEnd2 = pVect1 + 4 * qty4;
+
+        __m256 sum256 = _mm256_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            //_mm_prefetch((char*)(pVect2 + 16), _MM_HINT_T0);
+
+            __m256 v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            __m256 v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            sum256 = _mm256_add_ps(sum256, _mm256_mul_ps(v1, v2));
+
+            v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            sum256 = _mm256_add_ps(sum256, _mm256_mul_ps(v1, v2));
+        }
+
+        __m128 v1, v2;
+        __m128 sum_prod = _mm_add_ps(_mm256_extractf128_ps(sum256, 0), _mm256_extractf128_ps(sum256, 1));
+
+        while (pVect1 < pEnd2) {
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+        }
+
+        _mm_store_ps(TmpRes, sum_prod);
+        float sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3];;
+        return 1.0f - sum;
+}
+
+#elif defined(USE_SSE)
+
+    static float
+    InnerProductSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float PORTABLE_ALIGN32 TmpRes[8];
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+        size_t qty16 = qty / 16;
+        size_t qty4 = qty / 4;
+
+        const float *pEnd1 = pVect1 + 16 * qty16;
+        const float *pEnd2 = pVect1 + 4 * qty4;
+
+        __m128 v1, v2;
+        __m128 sum_prod = _mm_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+        }
+
+        while (pVect1 < pEnd2) {
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+        }
+
+        _mm_store_ps(TmpRes, sum_prod);
+        float sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3];
+
+        return 1.0f - sum;
+    }
+
+#endif
+
+#if defined(USE_AVX)
+
+    static float
+    InnerProductSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float PORTABLE_ALIGN32 TmpRes[8];
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+        size_t qty16 = qty / 16;
+
+
+        const float *pEnd1 = pVect1 + 16 * qty16;
+
+        __m256 sum256 = _mm256_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            //_mm_prefetch((char*)(pVect2 + 16), _MM_HINT_T0);
+
+            __m256 v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            __m256 v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            sum256 = _mm256_add_ps(sum256, _mm256_mul_ps(v1, v2));
+
+            v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            sum256 = _mm256_add_ps(sum256, _mm256_mul_ps(v1, v2));
+        }
+
+        _mm256_store_ps(TmpRes, sum256);
+        float sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7];
+
+        return 1.0f - sum;
+    }
+
+#elif defined(USE_SSE)
+
+      static float
+      InnerProductSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float PORTABLE_ALIGN32 TmpRes[8];
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+        size_t qty16 = qty / 16;
+
+        const float *pEnd1 = pVect1 + 16 * qty16;
+
+        __m128 v1, v2;
+        __m128 sum_prod = _mm_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            sum_prod = _mm_add_ps(sum_prod, _mm_mul_ps(v1, v2));
+        }
+        _mm_store_ps(TmpRes, sum_prod);
+        float sum = TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3];
+
+        return 1.0f - sum;
+    }
+
+#endif
+
+#if defined(USE_SSE) || defined(USE_AVX)
+    static float
+    InnerProductSIMD16ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        size_t qty = *((size_t *) qty_ptr);
+        size_t qty16 = qty >> 4 << 4;
+        float res = InnerProductSIMD16Ext(pVect1v, pVect2v, &qty16);
+        float *pVect1 = (float *) pVect1v + qty16;
+        float *pVect2 = (float *) pVect2v + qty16;
+
+        size_t qty_left = qty - qty16;
+        float res_tail = InnerProduct(pVect1, pVect2, &qty_left);
+        return res + res_tail - 1.0f;
+    }
+
+    static float
+    InnerProductSIMD4ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        size_t qty = *((size_t *) qty_ptr);
+        size_t qty4 = qty >> 2 << 2;
+
+        float res = InnerProductSIMD4Ext(pVect1v, pVect2v, &qty4);
+        size_t qty_left = qty - qty4;
+
+        float *pVect1 = (float *) pVect1v + qty4;
+        float *pVect2 = (float *) pVect2v + qty4;
+        float res_tail = InnerProduct(pVect1, pVect2, &qty_left);
+
+        return res + res_tail - 1.0f;
+    }
+#endif
+
+    class InnerProductSpace : public SpaceInterface<float> {
+
+        DISTFUNC<float> fstdistfunc_;
+        size_t data_size_;
+        size_t dim_;
+    public:
+        InnerProductSpace(size_t dim) {
+            fstdistfunc_ = InnerProduct;
+    #if defined(USE_AVX) || defined(USE_SSE)
+            if (dim % 16 == 0)
+                fstdistfunc_ = InnerProductSIMD16Ext;
+            else if (dim % 4 == 0)
+                fstdistfunc_ = InnerProductSIMD4Ext;
+            else if (dim > 16)
+                fstdistfunc_ = InnerProductSIMD16ExtResiduals;
+            else if (dim > 4)
+                fstdistfunc_ = InnerProductSIMD4ExtResiduals;
+    #endif
+            dim_ = dim;
+            data_size_ = dim * sizeof(float);
+        }
+
+        size_t get_data_size() {
+            return data_size_;
+        }
+
+        DISTFUNC<float> get_dist_func() {
+            return fstdistfunc_;
+        }
+
+        void *get_dist_func_param() {
+            return &dim_;
+        }
+
+    ~InnerProductSpace() {}
+    };
+
+
+}

data/ext/hnswlib/src/space_l2.h

@@ -0,0 +1,281 @@
+#pragma once
+#include "hnswlib.h"
+
+namespace hnswlib {
+
+    static float
+    L2Sqr(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+        float res = 0;
+        for (size_t i = 0; i < qty; i++) {
+            float t = *pVect1 - *pVect2;
+            pVect1++;
+            pVect2++;
+            res += t * t;
+        }
+        return (res);
+    }
+
+#if defined(USE_AVX)
+
+    // Favor using AVX if available.
+    static float
+    L2SqrSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+        float PORTABLE_ALIGN32 TmpRes[8];
+        size_t qty16 = qty >> 4;
+
+        const float *pEnd1 = pVect1 + (qty16 << 4);
+
+        __m256 diff, v1, v2;
+        __m256 sum = _mm256_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            diff = _mm256_sub_ps(v1, v2);
+            sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
+
+            v1 = _mm256_loadu_ps(pVect1);
+            pVect1 += 8;
+            v2 = _mm256_loadu_ps(pVect2);
+            pVect2 += 8;
+            diff = _mm256_sub_ps(v1, v2);
+            sum = _mm256_add_ps(sum, _mm256_mul_ps(diff, diff));
+        }
+
+        _mm256_store_ps(TmpRes, sum);
+        return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3] + TmpRes[4] + TmpRes[5] + TmpRes[6] + TmpRes[7];
+    }
+
+#elif defined(USE_SSE)
+
+    static float
+    L2SqrSIMD16Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+        float PORTABLE_ALIGN32 TmpRes[8];
+        size_t qty16 = qty >> 4;
+
+        const float *pEnd1 = pVect1 + (qty16 << 4);
+
+        __m128 diff, v1, v2;
+        __m128 sum = _mm_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            //_mm_prefetch((char*)(pVect2 + 16), _MM_HINT_T0);
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            diff = _mm_sub_ps(v1, v2);
+            sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            diff = _mm_sub_ps(v1, v2);
+            sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            diff = _mm_sub_ps(v1, v2);
+            sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
+
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            diff = _mm_sub_ps(v1, v2);
+            sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
+        }
+
+        _mm_store_ps(TmpRes, sum);
+        return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3];
+    }
+#endif
+
+#if defined(USE_SSE) || defined(USE_AVX)
+    static float
+    L2SqrSIMD16ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        size_t qty = *((size_t *) qty_ptr);
+        size_t qty16 = qty >> 4 << 4;
+        float res = L2SqrSIMD16Ext(pVect1v, pVect2v, &qty16);
+        float *pVect1 = (float *) pVect1v + qty16;
+        float *pVect2 = (float *) pVect2v + qty16;
+
+        size_t qty_left = qty - qty16;
+        float res_tail = L2Sqr(pVect1, pVect2, &qty_left);
+        return (res + res_tail);
+    }
+#endif
+
+
+#ifdef USE_SSE
+    static float
+    L2SqrSIMD4Ext(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        float PORTABLE_ALIGN32 TmpRes[8];
+        float *pVect1 = (float *) pVect1v;
+        float *pVect2 = (float *) pVect2v;
+        size_t qty = *((size_t *) qty_ptr);
+
+
+        size_t qty4 = qty >> 2;
+
+        const float *pEnd1 = pVect1 + (qty4 << 2);
+
+        __m128 diff, v1, v2;
+        __m128 sum = _mm_set1_ps(0);
+
+        while (pVect1 < pEnd1) {
+            v1 = _mm_loadu_ps(pVect1);
+            pVect1 += 4;
+            v2 = _mm_loadu_ps(pVect2);
+            pVect2 += 4;
+            diff = _mm_sub_ps(v1, v2);
+            sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
+        }
+        _mm_store_ps(TmpRes, sum);
+        return TmpRes[0] + TmpRes[1] + TmpRes[2] + TmpRes[3];
+    }
+
+    static float
+    L2SqrSIMD4ExtResiduals(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
+        size_t qty = *((size_t *) qty_ptr);
+        size_t qty4 = qty >> 2 << 2;
+
+        float res = L2SqrSIMD4Ext(pVect1v, pVect2v, &qty4);
+        size_t qty_left = qty - qty4;
+
+        float *pVect1 = (float *) pVect1v + qty4;
+        float *pVect2 = (float *) pVect2v + qty4;
+        float res_tail = L2Sqr(pVect1, pVect2, &qty_left);
+
+        return (res + res_tail);
+    }
+#endif
+
+    class L2Space : public SpaceInterface<float> {
+
+        DISTFUNC<float> fstdistfunc_;
+        size_t data_size_;
+        size_t dim_;
+    public:
+        L2Space(size_t dim) {
+            fstdistfunc_ = L2Sqr;
+        #if defined(USE_SSE) || defined(USE_AVX)
+            if (dim % 16 == 0)
+                fstdistfunc_ = L2SqrSIMD16Ext;
+            else if (dim % 4 == 0)
+                fstdistfunc_ = L2SqrSIMD4Ext;
+            else if (dim > 16)
+                fstdistfunc_ = L2SqrSIMD16ExtResiduals;
+            else if (dim > 4)
+                fstdistfunc_ = L2SqrSIMD4ExtResiduals;
+        #endif
+            dim_ = dim;
+            data_size_ = dim * sizeof(float);
+        }
+
+        size_t get_data_size() {
+            return data_size_;
+        }
+
+        DISTFUNC<float> get_dist_func() {
+            return fstdistfunc_;
+        }
+
+        void *get_dist_func_param() {
+            return &dim_;
+        }
+
+        ~L2Space() {}
+    };
+
+    static int
+    L2SqrI4x(const void *__restrict pVect1, const void *__restrict pVect2, const void *__restrict qty_ptr) {
+
+        size_t qty = *((size_t *) qty_ptr);
+        int res = 0;
+        unsigned char *a = (unsigned char *) pVect1;
+        unsigned char *b = (unsigned char *) pVect2;
+
+        qty = qty >> 2;
+        for (size_t i = 0; i < qty; i++) {
+
+            res += ((*a) - (*b)) * ((*a) - (*b));
+            a++;
+            b++;
+            res += ((*a) - (*b)) * ((*a) - (*b));
+            a++;
+            b++;
+            res += ((*a) - (*b)) * ((*a) - (*b));
+            a++;
+            b++;
+            res += ((*a) - (*b)) * ((*a) - (*b));
+            a++;
+            b++;
+        }
+        return (res);
+    }
+
+    static int L2SqrI(const void* __restrict pVect1, const void* __restrict pVect2, const void* __restrict qty_ptr) {
+        size_t qty = *((size_t*)qty_ptr);
+        int res = 0;
+        unsigned char* a = (unsigned char*)pVect1;
+        unsigned char* b = (unsigned char*)pVect2;
+
+        for(size_t i = 0; i < qty; i++)
+        {
+            res += ((*a) - (*b)) * ((*a) - (*b));
+            a++;
+            b++;
+        }
+        return (res);
+    }
+
+    class L2SpaceI : public SpaceInterface<int> {
+
+        DISTFUNC<int> fstdistfunc_;
+        size_t data_size_;
+        size_t dim_;
+    public:
+        L2SpaceI(size_t dim) {
+            if(dim % 4 == 0) {
+                fstdistfunc_ = L2SqrI4x;
+            }
+            else {
+                fstdistfunc_ = L2SqrI;
+            }
+            dim_ = dim;
+            data_size_ = dim * sizeof(unsigned char);
+        }
+
+        size_t get_data_size() {
+            return data_size_;
+        }
+
+        DISTFUNC<int> get_dist_func() {
+            return fstdistfunc_;
+        }
+
+        void *get_dist_func_param() {
+            return &dim_;
+        }
+
+        ~L2SpaceI() {}
+    };
+
+
+}