Fix inner product distance code to vectorize on arm (affects all arch)

examples/cpp/example_search.cpp

@@ -1,58 +1,120 @@
 #include "../../hnswlib/hnswlib.h"
 
+#include <chrono>
 
-int main() {
-    int dim = 16;               // Dimension of the elements
-    int max_elements = 10000;   // Maximum number of elements, should be known beforehand
-    int M = 16;                 // Tightly connected with internal dimensionality of the data
-                                // strongly affects the memory consumption
-    int ef_construction = 200;  // Controls index search speed/build speed tradeoff
+void normalize_data_for_inner_product(float* data, int dim, int max_elements) {
+    for (int i = 0; i < max_elements; i++) {
+        double norm_sq = 0.0;
 
-    // Initing index
-    hnswlib::L2Space space(dim);
-    hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, max_elements, M, ef_construction);
+        for (int j = 0; j < dim; j++) {
+            float value = data[i * dim + j];
+            norm_sq += static_cast<double>(value) * value;
+        }
 
-    // Generate random data
-    std::mt19937 rng;
-    rng.seed(47);
-    std::uniform_real_distribution<> distrib_real;
-    float* data = new float[dim * max_elements];
-    for (int i = 0; i < dim * max_elements; i++) {
-        data[i] = distrib_real(rng);
+        double norm = std::sqrt(norm_sq);
+
+        if (norm > 0.0) {
+            for (int j = 0; j < dim; j++) {
+                data[i * dim + j] =
+                    static_cast<float>(data[i * dim + j] / norm);
+            }
+        }
     }
+}
+
+template <typename SpaceType>
+void run_test(const std::string& name, SpaceType& space, float* data, int dim, int max_elements) {
+    int M = 16;
+    int ef_construction = 200;
+
+    std::cout << "Running " << name << " test...\n";
+
+    // Initing index
+    hnswlib::HierarchicalNSW<float>* alg_hnsw =
+        new hnswlib::HierarchicalNSW<float>(&space, max_elements, M, ef_construction);
+
+    auto build_start = std::chrono::high_resolution_clock::now();
 
     // Add data to index
     for (int i = 0; i < max_elements; i++) {
         alg_hnsw->addPoint(data + i * dim, i);
     }
 
+    auto build_end = std::chrono::high_resolution_clock::now();
+
+    std::chrono::duration<double> build_seconds = build_end - build_start;
+    auto build_ms =
+        std::chrono::duration<double, std::milli>(build_end - build_start);
+
+    std::cout << name << " index build time: "
+              << build_seconds.count() << " seconds ("
+              << build_ms.count() << " ms)\n";
+
     // Query the elements for themselves and measure recall
     float correct = 0;
     for (int i = 0; i < max_elements; i++) {
-        std::priority_queue<std::pair<float, hnswlib::labeltype>> result = alg_hnsw->searchKnn(data + i * dim, 1);
+        std::priority_queue<std::pair<float, hnswlib::labeltype>> result =
+            alg_hnsw->searchKnn(data + i * dim, 1);
+
         hnswlib::labeltype label = result.top().second;
         if (label == i) correct++;
     }
+
     float recall = correct / max_elements;
-    std::cout << "Recall: " << recall << "\n";
+    std::cout << name << " Recall: " << recall << "\n";
 
     // Serialize index
-    std::string hnsw_path = "hnsw.bin";
+    std::string hnsw_path = name + "_hnsw.bin";
     alg_hnsw->saveIndex(hnsw_path);
     delete alg_hnsw;
 
     // Deserialize index and check recall
     alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, hnsw_path);
+
     correct = 0;
     for (int i = 0; i < max_elements; i++) {
-        std::priority_queue<std::pair<float, hnswlib::labeltype>> result = alg_hnsw->searchKnn(data + i * dim, 1);
+        std::priority_queue<std::pair<float, hnswlib::labeltype>> result =
+            alg_hnsw->searchKnn(data + i * dim, 1);
+
         hnswlib::labeltype label = result.top().second;
-        if (label == i) correct++;
+        if (label == i)  correct++;
     }
-    recall = (float)correct / max_elements;
-    std::cout << "Recall of deserialized index: " << recall << "\n";
 
-    delete[] data;
+    recall = correct / max_elements;
+    std::cout << name << " Recall of deserialized index: "
+              << recall << "\n\n";
+
     delete alg_hnsw;
+}
+
+int main() {
+    int dim = 16;
+    int max_elements = 10000;
+
+    // Generate random data
+    std::mt19937 rng;
+    rng.seed(47);
+    std::uniform_real_distribution<float> distrib_real(0.0f, 1.0f);
+
+    float* data_l2 = new float[dim * max_elements];
+    float* data_ip = new float[dim * max_elements];
+
+    for (int i = 0; i < dim * max_elements; i++) {
+        float value = distrib_real(rng);
+        data_l2[i] = value;
+        data_ip[i] = value;
+    }
+
+    hnswlib::L2Space l2_space(dim);
+    run_test("l2", l2_space, data_l2, dim, max_elements);
+
+    normalize_data_for_inner_product(data_ip, dim, max_elements);
+
+    hnswlib::InnerProductSpace ip_space(dim);
+    run_test("ip_normalized", ip_space, data_ip, dim, max_elements);
+
+    delete[] data_l2;
+    delete[] data_ip;
+
     return 0;
 }

hnswlib/space_ip.h

@@ -5,10 +5,13 @@ 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];
+    const float *vec1 = (const float *) pVect1;
+    const float *vec2 = (const float *) pVect2;
+    const size_t qty = *((const size_t *) qty_ptr);
+
+    float res = 0.0f;
+    for (size_t i = 0; i < qty; i++) {
+       res += vec1[i] * vec2[i];
     }
     return res;
 }