prefetching inside get_distance

Author: Ananya Sutradhar

scripts/analyze_dataset_span_simple.py

@@ -190,23 +190,6 @@ def main():
     print(f"{'Mean normalization':<25} {suggestions['mean_norm']:<15.6f} {'0.0':<15} {'Simple scaling only'}")
     print(f"{'Median normalization':<25} {suggestions['median_norm']:<15.6f} {'0.0':<15} {'Robust to outliers'}")
 
-    # Practical recommendations
-    print("\n" + "="*60)
-    print("PRACTICAL RECOMMENDATIONS")
-    print("="*60)
-    
-    print("For DiskANN filter + distance combination:")
-    print(f"1. RECOMMENDED: Min-Max normalization")
-    print(f"   - Apply: normalized_dist = (distance - {stats['min']:.6f}) * {suggestions['span_norm']:.6f}")
-    print(f"   - Result: distances in [0, 1] range, same as filter similarities")
-    
-    print(f"\n2. ALTERNATIVE: Robust normalization (if outliers present)")
-    print(f"   - Apply: normalized_dist = (distance - {stats['q25']:.6f}) * {suggestions['robust_norm']:.6f}")
-    print(f"   - Result: most distances in [0, 1], outliers may exceed 1")
-    
-    print(f"\n3. SIMPLE: Mean normalization (if centering not needed)")
-    print(f"   - Apply: normalized_dist = distance * {suggestions['mean_norm']:.6f}")
-    print(f"   - Result: average distance becomes 1.0")
 
     # Always save simple normalization factors to text file
     norm_config_file = args.norm_factors if args.norm_factors else 'normalization_factors.txt'

src/distance.cpp

@@ -120,15 +120,31 @@ float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t size) c
 {
 #ifdef _WINDOWS
 #ifdef USE_AVX2
+    // Prefetch the start of both vectors
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
     __m256 r = _mm256_setzero_ps();
     char *pX = (char *)a, *pY = (char *)b;
+    const char *original_pX = pX;
+    const char *original_pY = pY;
+    uint32_t prefetch_offset = 64; // Prefetch 64 bytes ahead
+    
     while (size >= 32)
     {
+        // Prefetch ahead for better cache performance
+        if (size > prefetch_offset) 
+        {
+            _mm_prefetch(original_pX + prefetch_offset, _MM_HINT_T0);
+            _mm_prefetch(original_pY + prefetch_offset, _MM_HINT_T0);
+        }
+        
         __m256i r1 = _mm256_subs_epi8(_mm256_loadu_si256((__m256i *)pX), _mm256_loadu_si256((__m256i *)pY));
         r = _mm256_add_ps(r, _mm256_mul_epi8(r1, r1));
         pX += 32;
         pY += 32;
         size -= 32;
+        prefetch_offset += 32;
     }
     while (size > 0)
     {
@@ -141,19 +157,39 @@ float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t size) c
     r = _mm256_hadd_ps(_mm256_hadd_ps(r, r), r);
     return r.m256_f32[0] + r.m256_f32[4];
 #else
+    // Prefetch the start of both vectors for non-AVX2 fallback
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
     int32_t result = 0;
 #pragma omp simd reduction(+ : result) aligned(a, b : 8)
     for (int32_t i = 0; i < (int32_t)size; i++)
     {
+        // Prefetch ahead every 64 bytes (64 int8_t values)
+        if (i % 64 == 0 && i + 64 < (int32_t)size) 
+        {
+            _mm_prefetch((const char*)(a + i + 64), _MM_HINT_T0);
+            _mm_prefetch((const char*)(b + i + 64), _MM_HINT_T0);
+        }
         result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i]));
     }
     return (float)result;
 #endif
 #else
+    // Prefetch the start of both vectors for Linux version
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
     int32_t result = 0;
 #pragma omp simd reduction(+ : result) aligned(a, b : 8)
     for (int32_t i = 0; i < (int32_t)size; i++)
     {
+        // Prefetch ahead every 64 bytes (64 int8_t values)
+        if (i % 64 == 0 && i + 64 < (int32_t)size) 
+        {
+            _mm_prefetch((const char*)(a + i + 64), _MM_HINT_T0);
+            _mm_prefetch((const char*)(b + i + 64), _MM_HINT_T0);
+        }
         result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i]));
     }
     return (float)result;
@@ -162,12 +198,22 @@ float DistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t size) c
 
 float DistanceL2UInt8::compare(const uint8_t *a, const uint8_t *b, uint32_t size) const
 {
+    // Prefetch the start of both vectors
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
     uint32_t result = 0;
 #ifndef _WINDOWS
 #pragma omp simd reduction(+ : result) aligned(a, b : 8)
 #endif
     for (int32_t i = 0; i < (int32_t)size; i++)
     {
+        // Prefetch ahead every 64 bytes (64 uint8_t values)
+        if (i % 64 == 0 && i + 64 < (int32_t)size) 
+        {
+            _mm_prefetch((const char*)(a + i + 64), _MM_HINT_T0);
+            _mm_prefetch((const char*)(b + i + 64), _MM_HINT_T0);
+        }
         result += ((int32_t)((int16_t)a[i] - (int16_t)b[i])) * ((int32_t)((int16_t)a[i] - (int16_t)b[i]));
     }
     return (float)result;
@@ -209,11 +255,21 @@ float DistanceL2Float::compare(const float *a, const float *b, uint32_t size) co
     // horizontal add sum
     result = _mm256_reduce_add_ps(sum);
 #else
+    // Prefetch the start of both vectors for non-AVX2 fallback
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
 #ifndef _WINDOWS
 #pragma omp simd reduction(+ : result) aligned(a, b : 32)
 #endif
     for (int32_t i = 0; i < (int32_t)size; i++)
     {
+        // Prefetch ahead every 16 floats (64 bytes)
+        if (i % 16 == 0 && i + 16 < (int32_t)size) 
+        {
+            _mm_prefetch((const char*)(a + i + 16), _MM_HINT_T0);
+            _mm_prefetch((const char*)(b + i + 16), _MM_HINT_T0);
+        }
         result += (a[i] - b[i]) * (a[i] - b[i]);
     }
 #endif
@@ -271,18 +327,34 @@ float AVXDistanceL2Int8::compare(const int8_t *a, const int8_t *b, uint32_t leng
 
 float AVXDistanceL2Float::compare(const float *a, const float *b, uint32_t length) const
 {
+    // Prefetch the start of both vectors
+    _mm_prefetch((const char*)a, _MM_HINT_T0);
+    _mm_prefetch((const char*)b, _MM_HINT_T0);
+    
     __m128 diff, v1, v2;
     __m128 sum = _mm_set1_ps(0);
+    
+    const float *original_a = a;
+    const float *original_b = b;
+    uint32_t prefetch_offset = 64; // Prefetch 64 bytes ahead (16 floats)
 
     while (length >= 4)
     {
+        // Prefetch ahead for better cache performance
+        if (length > prefetch_offset) 
+        {
+            _mm_prefetch((const char*)(original_a + prefetch_offset), _MM_HINT_T0);
+            _mm_prefetch((const char*)(original_b + prefetch_offset), _MM_HINT_T0);
+        }
+        
         v1 = _mm_loadu_ps(a);
         a += 4;
         v2 = _mm_loadu_ps(b);
         b += 4;
         diff = _mm_sub_ps(v1, v2);
         sum = _mm_add_ps(sum, _mm_mul_ps(diff, diff));
         length -= 4;
+        prefetch_offset += 4;
     }
 
     return sum.m128_f32[0] + sum.m128_f32[1] + sum.m128_f32[2] + sum.m128_f32[3];

src/in_mem_data_store.cpp

@@ -7,6 +7,12 @@
 
 #include "utils.h"
 
+#ifdef _WIN32
+#include <intrin.h>
+#else
+#include <x86intrin.h>
+#endif
+
 namespace diskann
 {
 
@@ -197,17 +203,38 @@ void InMemDataStore<data_t>::preprocess_query(const data_t *query, AbstractScrat
 
 template <typename data_t> float InMemDataStore<data_t>::get_distance(const data_t *query, const location_t loc) const
 {
-    return _distance_fn->compare(query, _data + _aligned_dim * loc, (uint32_t)_aligned_dim);
+    // Prefetch query vector for better cache performance
+    _mm_prefetch((const char*)query, _MM_HINT_T0);
+    
+    // Prefetch the current vector data
+    const data_t *vector_data = _data + _aligned_dim * loc;
+    _mm_prefetch((const char*)vector_data, _MM_HINT_T0);
+    
+    return _distance_fn->compare(query, vector_data, (uint32_t)_aligned_dim);
 }
 
 template <typename data_t>
 void InMemDataStore<data_t>::get_distance(const data_t *query, const location_t *locations,
                                           const uint32_t location_count, float *distances,
                                           AbstractScratch<data_t> *scratch_space) const
 {
+    // Prefetch query vector once for the entire batch
+    _mm_prefetch((const char*)query, _MM_HINT_T0);
+    
     for (location_t i = 0; i < location_count; i++)
     {
-        distances[i] = _distance_fn->compare(query, _data + locations[i] * _aligned_dim, (uint32_t)this->_aligned_dim);
+        // Prefetch current vector data
+        const data_t *current_vector = _data + locations[i] * _aligned_dim;
+        _mm_prefetch((const char*)current_vector, _MM_HINT_T0);
+        
+        // Prefetch next vector data if available
+        if (i + 1 < location_count)
+        {
+            const data_t *next_vector = _data + locations[i + 1] * _aligned_dim;
+            _mm_prefetch((const char*)next_vector, _MM_HINT_T0);
+        }
+        
+        distances[i] = _distance_fn->compare(query, current_vector, (uint32_t)this->_aligned_dim);
     }
 }
 
@@ -222,10 +249,24 @@ template <typename data_t>
 void InMemDataStore<data_t>::get_distance(const data_t *preprocessed_query, const std::vector<location_t> &ids,
                                           std::vector<float> &distances, AbstractScratch<data_t> *scratch_space) const
 {
+    // Prefetch query vector once for the entire batch
+    _mm_prefetch((const char*)preprocessed_query, _MM_HINT_T0);
+    
     for (int i = 0; i < ids.size(); i++)
     {
+        // Prefetch current vector data
+        const data_t *current_vector = _data + ids[i] * _aligned_dim;
+        _mm_prefetch((const char*)current_vector, _MM_HINT_T0);
+        
+        // Prefetch next vector data if available
+        if (i + 1 < (int)ids.size())
+        {
+            const data_t *next_vector = _data + ids[i + 1] * _aligned_dim;
+            _mm_prefetch((const char*)next_vector, _MM_HINT_T0);
+        }
+        
         distances[i] =
-            _distance_fn->compare(preprocessed_query, _data + ids[i] * _aligned_dim, (uint32_t)this->_aligned_dim);
+            _distance_fn->compare(preprocessed_query, current_vector, (uint32_t)this->_aligned_dim);
     }
 }
 

src/index.cpp

@@ -1074,7 +1074,8 @@ inline void Index<T, TagT, LabelT>::calculate_jaccard_similarity_batch(
 
                 // Check if any label in this clause matches the point
                 for (const auto& label : clause) {
-                    if (_labels_to_points_set[label].count(point_id)) {
+                    auto it = _labels_to_points_set.find(label);
+                    if (it != _labels_to_points_set.end() && it->second.count(point_id)) {
                         matching_clauses++;
                         break; // Found match, move to next clause
                     }

src/pq_data_store.cpp

@@ -6,6 +6,12 @@
 #include "utils.h"
 #include "distance.h"
 
+#ifdef _WIN32
+#include <intrin.h>
+#else
+#include <x86intrin.h>
+#endif
+
 namespace diskann
 {
 
@@ -180,6 +186,17 @@ void PQDataStore<data_t>::get_distance(const data_t *preprocessed_query, const l
     {
         throw diskann::ANNException("PQScratch not set in scratch space.", -1);
     }
+    
+    // Prefetch preprocessed query data
+    _mm_prefetch((const char*)preprocessed_query, _MM_HINT_T0);
+    
+    // Prefetch some of the quantized data for the locations we'll process
+    for (uint32_t i = 0; i < location_count && i < 4; i++)
+    {
+        const uint8_t *quantized_loc = _quantized_data + locations[i] * this->_num_chunks;
+        _mm_prefetch((const char*)quantized_loc, _MM_HINT_T0);
+    }
+    
     diskann::aggregate_coords(locations, location_count, _quantized_data, this->_num_chunks,
                               pq_scratch->aligned_pq_coord_scratch);
     _pq_distance_fn->preprocessed_distance(*pq_scratch, location_count, distances);
@@ -198,6 +215,18 @@ void PQDataStore<data_t>::get_distance(const data_t *preprocessed_query, const s
     {
         throw diskann::ANNException("PQScratch not set in scratch space.", -1);
     }
+    
+    // Prefetch preprocessed query data
+    _mm_prefetch((const char*)preprocessed_query, _MM_HINT_T0);
+    
+    // Prefetch some of the quantized data for the locations we'll process
+    size_t prefetch_count = ids.size() < 4 ? ids.size() : 4;
+    for (size_t i = 0; i < prefetch_count; i++)
+    {
+        const uint8_t *quantized_loc = _quantized_data + ids[i] * this->_num_chunks;
+        _mm_prefetch((const char*)quantized_loc, _MM_HINT_T0);
+    }
+    
     diskann::aggregate_coords(ids, _quantized_data, this->_num_chunks, pq_scratch->aligned_pq_coord_scratch);
     _pq_distance_fn->preprocessed_distance(*pq_scratch, (location_t)ids.size(), distances);
 }