[MOD-8200] [MOD-8202] INT8 index

src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h

@@ -57,6 +57,8 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
 
 friend class BF16TieredTest;
 friend class FP16TieredTest;
+friend class INT8TieredTest;
+friend class CommonTypeMetricTieredTests_TestDataSizeTieredHNSW_Test;
 
 INDEX_TEST_FRIEND_CLASS(BM_VecSimBasics)
 INDEX_TEST_FRIEND_CLASS(BM_VecSimCommon)

src/VecSim/index_factories/brute_force_factory.cpp

@@ -33,10 +33,12 @@
 static AbstractIndexInitParams NewAbstractInitParams(const VecSimParams *params) {
 
     const BFParams *bfParams = &params->algoParams.bfParams;
+    size_t dataSize = VecSimParams_GetDataSize(bfParams->type, bfParams->dim, bfParams->metric);
     AbstractIndexInitParams abstractInitParams = {.allocator =
                                                       VecSimAllocator::newVecsimAllocator(),
                                                   .dim = bfParams->dim,
                                                   .vecType = bfParams->type,
+                                                  .dataSize = dataSize,
                                                   .metric = bfParams->metric,
                                                   .blockSize = bfParams->blockSize,
                                                   .multi = bfParams->multi,
@@ -52,32 +54,30 @@
 
 VecSimIndex *NewIndex(const BFParams *bfparams, const AbstractIndexInitParams &abstractInitParams,
                       bool is_normalized) {
-    // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors
-    // and query blobs.
-    VecSimMetric metric;
-    if (is_normalized && bfparams->metric == VecSimMetric_Cosine) {
-        metric = VecSimMetric_IP;
-    } else {
-        metric = bfparams->metric;
-    }
+
     if (bfparams->type == VecSimType_FLOAT32) {
         IndexComponents<float, float> indexComponents = CreateIndexComponents<float, float>(
-            abstractInitParams.allocator, metric, bfparams->dim);
+            abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float>(bfparams, abstractInitParams, indexComponents);
     } else if (bfparams->type == VecSimType_FLOAT64) {
         IndexComponents<double, double> indexComponents = CreateIndexComponents<double, double>(
-            abstractInitParams.allocator, metric, bfparams->dim);
+            abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<double>(bfparams, abstractInitParams, indexComponents);
     } else if (bfparams->type == VecSimType_BFLOAT16) {
         IndexComponents<bfloat16, float> indexComponents = CreateIndexComponents<bfloat16, float>(
-            abstractInitParams.allocator, metric, bfparams->dim);
+            abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<bfloat16, float>(bfparams, abstractInitParams,
                                                              indexComponents);
     } else if (bfparams->type == VecSimType_FLOAT16) {
         IndexComponents<float16, float> indexComponents = CreateIndexComponents<float16, float>(
-            abstractInitParams.allocator, metric, bfparams->dim);
+            abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float16, float>(bfparams, abstractInitParams,
                                                             indexComponents);
+    } else if (bfparams->type == VecSimType_INT8) {
+        IndexComponents<int8_t, float> indexComponents = CreateIndexComponents<int8_t, float>(
+            abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized);
+        return NewIndex_ChooseMultiOrSingle<int8_t, float>(bfparams, abstractInitParams,
+                                                           indexComponents);
     }
 
     // If we got here something is wrong.
@@ -117,6 +117,11 @@
     } else if (params->type == VecSimType_FLOAT16) {
         est += EstimateComponentsMemory<float16, float>(params->metric, is_normalized);
         est += EstimateInitialSize_ChooseMultiOrSingle<float16, float>(params->multi);
+    } else if (params->type == VecSimType_INT8) {
+        est += EstimateComponentsMemory<int8_t, float>(params->metric, is_normalized);
+        est += EstimateInitialSize_ChooseMultiOrSingle<int8_t, float>(params->multi);
+    } else {
+        throw std::invalid_argument("Invalid params->type");
     }
 
     est += sizeof(DataBlocksContainer) + allocations_overhead;

src/VecSim/index_factories/components/components_factory.h

@@ -14,14 +14,24 @@
 
 template <typename DataType, typename DistType>
 IndexComponents<DataType, DistType>
-CreateIndexComponents(std::shared_ptr<VecSimAllocator> allocator, VecSimMetric metric, size_t dim) {
+CreateIndexComponents(std::shared_ptr<VecSimAllocator> allocator, VecSimMetric metric, size_t dim,
+                      bool is_normalized) {
     unsigned char alignment = 0;
     spaces::dist_func_t<DistType> distFunc =
         spaces::GetDistFunc<DataType, DistType>(metric, dim, &alignment);
     // Currently we have only one distance calculator implementation
     auto indexCalculator = new (allocator) DistanceCalculatorCommon<DistType>(allocator, distFunc);
 
-    PreprocessorsContainerParams ppParams = {.metric = metric, .dim = dim, .alignment = alignment};
+    // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors
+    // and query blobs.
+    VecSimMetric pp_metric;
+    if (is_normalized && metric == VecSimMetric_Cosine) {
+        pp_metric = VecSimMetric_IP;
+    } else {
+        pp_metric = metric;
+    }
+    PreprocessorsContainerParams ppParams = {
+        .metric = pp_metric, .dim = dim, .alignment = alignment};
     auto preprocessors = CreatePreprocessorsContainer<DataType>(allocator, ppParams);
 
     return {indexCalculator, preprocessors};

src/VecSim/index_factories/hnsw_factory.cpp

@@ -33,10 +33,14 @@
 
 static AbstractIndexInitParams NewAbstractInitParams(const VecSimParams *params) {
     const HNSWParams *hnswParams = &params->algoParams.hnswParams;
+
+    size_t dataSize =
+        VecSimParams_GetDataSize(hnswParams->type, hnswParams->dim, hnswParams->metric);
     AbstractIndexInitParams abstractInitParams = {.allocator =
                                                       VecSimAllocator::newVecsimAllocator(),
                                                   .dim = hnswParams->dim,
                                                   .vecType = hnswParams->type,
+                                                  .dataSize = dataSize,
                                                   .metric = hnswParams->metric,
                                                   .blockSize = hnswParams->blockSize,
                                                   .multi = hnswParams->multi,
@@ -48,36 +52,32 @@
     const HNSWParams *hnswParams = &params->algoParams.hnswParams;
     AbstractIndexInitParams abstractInitParams = NewAbstractInitParams(params);
 
-    // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors
-    // and query blobs.
-    VecSimMetric metric;
-    if (is_normalized && hnswParams->metric == VecSimMetric_Cosine) {
-        metric = VecSimMetric_IP;
-    } else {
-        metric = hnswParams->metric;
-    }
-
     if (hnswParams->type == VecSimType_FLOAT32) {
         IndexComponents<float, float> indexComponents = CreateIndexComponents<float, float>(
-            abstractInitParams.allocator, metric, hnswParams->dim);
+            abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float>(hnswParams, abstractInitParams, indexComponents);
 
     } else if (hnswParams->type == VecSimType_FLOAT64) {
         IndexComponents<double, double> indexComponents = CreateIndexComponents<double, double>(
-            abstractInitParams.allocator, metric, hnswParams->dim);
+            abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<double>(hnswParams, abstractInitParams,
                                                     indexComponents);
 
     } else if (hnswParams->type == VecSimType_BFLOAT16) {
         IndexComponents<bfloat16, float> indexComponents = CreateIndexComponents<bfloat16, float>(
-            abstractInitParams.allocator, metric, hnswParams->dim);
+            abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<bfloat16, float>(hnswParams, abstractInitParams,
                                                              indexComponents);
     } else if (hnswParams->type == VecSimType_FLOAT16) {
         IndexComponents<float16, float> indexComponents = CreateIndexComponents<float16, float>(
-            abstractInitParams.allocator, metric, hnswParams->dim);
+            abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float16, float>(hnswParams, abstractInitParams,
                                                             indexComponents);
+    } else if (hnswParams->type == VecSimType_INT8) {
+        IndexComponents<int8_t, float> indexComponents = CreateIndexComponents<int8_t, float>(
+            abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
+        return NewIndex_ChooseMultiOrSingle<int8_t, float>(hnswParams, abstractInitParams,
+                                                           indexComponents);
     }
 
     // If we got here something is wrong.
@@ -114,6 +114,11 @@
     } else if (params->type == VecSimType_FLOAT16) {
         est += EstimateComponentsMemory<float16, float>(params->metric, is_normalized);
         est += EstimateInitialSize_ChooseMultiOrSingle<float16, float>(params->multi);
+    } else if (params->type == VecSimType_INT8) {
+        est += EstimateComponentsMemory<int8_t, float>(params->metric, is_normalized);
+        est += EstimateInitialSize_ChooseMultiOrSingle<int8_t, float>(params->multi);
+    } else {
+        throw std::invalid_argument("Invalid params->type");
     }
     return est;
 }
@@ -203,34 +208,32 @@
     VecSimParams vecsimParams = {.algo = VecSimAlgo_HNSWLIB,
                                  .algoParams = {.hnswParams = HNSWParams{params}}};
 
-    VecSimMetric metric;
-    if (is_normalized && params.metric == VecSimMetric_Cosine) {
-        metric = VecSimMetric_IP;
-    } else {
-        metric = params.metric;
-    }
-
     AbstractIndexInitParams abstractInitParams = NewAbstractInitParams(&vecsimParams);
     if (params.type == VecSimType_FLOAT32) {
         IndexComponents<float, float> indexComponents = CreateIndexComponents<float, float>(
-            abstractInitParams.allocator, metric, abstractInitParams.dim);
+            abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float>(input, &params, abstractInitParams,
                                                    indexComponents, version);
     } else if (params.type == VecSimType_FLOAT64) {
         IndexComponents<double, double> indexComponents = CreateIndexComponents<double, double>(
-            abstractInitParams.allocator, metric, abstractInitParams.dim);
+            abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<double>(input, &params, abstractInitParams,
                                                     indexComponents, version);
     } else if (params.type == VecSimType_BFLOAT16) {
         IndexComponents<bfloat16, float> indexComponents = CreateIndexComponents<bfloat16, float>(
-            abstractInitParams.allocator, metric, abstractInitParams.dim);
+            abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<bfloat16, float>(input, &params, abstractInitParams,
                                                              indexComponents, version);
     } else if (params.type == VecSimType_FLOAT16) {
         IndexComponents<float16, float> indexComponents = CreateIndexComponents<float16, float>(
-            abstractInitParams.allocator, metric, abstractInitParams.dim);
+            abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized);
         return NewIndex_ChooseMultiOrSingle<float16, float>(input, &params, abstractInitParams,
                                                             indexComponents, version);
+    } else if (params.type == VecSimType_INT8) {
+        IndexComponents<int8_t, float> indexComponents = CreateIndexComponents<int8_t, float>(
+            abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized);
+        return NewIndex_ChooseMultiOrSingle<int8_t, float>(input, &params, abstractInitParams,
+                                                           indexComponents, version);
     } else {
         auto bad_name = VecSimType_ToString(params.type);
         if (bad_name == nullptr) {

src/VecSim/index_factories/tiered_factory.cpp

@@ -42,9 +42,12 @@
     BFParams bf_params = NewBFParams(params);
 
     std::shared_ptr<VecSimAllocator> flat_allocator = VecSimAllocator::newVecsimAllocator();
+    size_t dataSize = VecSimParams_GetDataSize(bf_params.type, bf_params.dim, bf_params.metric);
+
     AbstractIndexInitParams abstractInitParams = {.allocator = flat_allocator,
                                                   .dim = bf_params.dim,
                                                   .vecType = bf_params.type,
+                                                  .dataSize = dataSize,
                                                   .metric = bf_params.metric,
                                                   .blockSize = bf_params.blockSize,
                                                   .multi = bf_params.multi,
@@ -80,6 +83,10 @@
         est += sizeof(TieredHNSWIndex<bfloat16, float>);
     } else if (hnsw_params.type == VecSimType_FLOAT16) {
         est += sizeof(TieredHNSWIndex<float16, float>);
+    } else if (hnsw_params.type == VecSimType_INT8) {
+        est += sizeof(TieredHNSWIndex<int8_t, float>);
+    } else {
+        throw std::invalid_argument("Invalid hnsw_params.type");
     }
 
     return est;
@@ -96,6 +103,8 @@
         return TieredHNSWFactory::NewIndex<bfloat16, float>(params);
     } else if (type == VecSimType_FLOAT16) {
         return TieredHNSWFactory::NewIndex<float16, float>(params);
+    } else if (type == VecSimType_INT8) {
+        return TieredHNSWFactory::NewIndex<int8_t, float>(params);
     }
     return nullptr; // Invalid type.
 }

src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_INT8.h

@@ -34,7 +34,7 @@ static inline int INT8_InnerProductImp(const void *pVect1v, const void *pVect2v,
     // Deal with remainder first. `dim` is more than 32, so we have at least one 32-int_8 block,
     // so mask loading is guaranteed to be safe
     if constexpr (residual % 32) {
-        __mmask32 mask = (1LU << (residual % 32)) - 1;
+        constexpr __mmask32 mask = (1LU << (residual % 32)) - 1;
         __m256i temp_a = _mm256_maskz_loadu_epi8(mask, pVect1);
         __m512i va = _mm512_cvtepi8_epi16(temp_a);
         pVect1 += residual % 32;

src/VecSim/spaces/normalize/compute_norm.h

@@ -0,0 +1,25 @@
+/*
+ *Copyright Redis Ltd. 2021 - present
+ *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or
+ *the Server Side Public License v1 (SSPLv1).
+ */
+
+#pragma once
+
+#include <cmath>
+
+namespace spaces {
+
+template <typename DataType>
+static inline float IntegralType_ComputeNorm(const DataType *vec, const size_t dim) {
+    int sum = 0;
+
+    for (size_t i = 0; i < dim; i++) {
+        // No need to cast to int because c++ integer promotion ensures vec[i] is promoted to int
+        // before multiplication.
+        sum += vec[i] * vec[i];
+    }
+    return sqrt(sum);
+}
+
+} // namespace spaces

src/VecSim/spaces/normalize/normalize_naive.h

@@ -8,6 +8,7 @@
 
 #include "VecSim/types/bfloat16.h"
 #include "VecSim/types/float16.h"
+#include "compute_norm.h"
 #include <cmath>
 #include <vector>
 
@@ -73,4 +74,13 @@ static inline void float16_normalizeVector(void *vec, const size_t dim) {
     }
 }
 
+static inline void int8_normalizeVector(void *vec, const size_t dim) {
+    int8_t *input_vector = static_cast<int8_t *>(vec);
+
+    float norm = IntegralType_ComputeNorm<int8_t>(input_vector, dim);
+
+    // Store norm at the end of the vector.
+    *reinterpret_cast<float *>(input_vector + dim) = norm;
+}
+
 } // namespace spaces

src/VecSim/spaces/spaces.cpp

@@ -108,4 +108,10 @@ normalizeVector_f<vecsim_types::float16> GetNormalizeFunc<vecsim_types::float16>
     return float16_normalizeVector;
 }
 
+/** The returned function computes the norm and stores it at the end of the given vector */
+template <>
+normalizeVector_f<int8_t> GetNormalizeFunc<int8_t>(void) {
+    return int8_normalizeVector;
+}
+
 } // namespace spaces

src/VecSim/utils/vec_utils.cpp

@@ -27,6 +27,7 @@ const char *VecSimCommonStrings::FLOAT32_STRING = "FLOAT32";
 const char *VecSimCommonStrings::FLOAT64_STRING = "FLOAT64";
 const char *VecSimCommonStrings::BFLOAT16_STRING = "BFLOAT16";
 const char *VecSimCommonStrings::FLOAT16_STRING = "FLOAT16";
+const char *VecSimCommonStrings::INT8_STRING = "INT8";
 const char *VecSimCommonStrings::INT32_STRING = "INT32";
 const char *VecSimCommonStrings::INT64_STRING = "INT64";
 
@@ -147,6 +148,8 @@ const char *VecSimType_ToString(VecSimType vecsimType) {
         return VecSimCommonStrings::BFLOAT16_STRING;
     case VecSimType_FLOAT16:
         return VecSimCommonStrings::FLOAT16_STRING;
+    case VecSimType_INT8:
+        return VecSimCommonStrings::INT8_STRING;
     case VecSimType_INT32:
         return VecSimCommonStrings::INT32_STRING;
     case VecSimType_INT64:
@@ -195,10 +198,20 @@ size_t VecSimType_sizeof(VecSimType type) {
         return sizeof(bfloat16);
     case VecSimType_FLOAT16:
         return sizeof(float16);
+    case VecSimType_INT8:
+        return sizeof(int8_t);
     case VecSimType_INT32:
         return sizeof(int32_t);
     case VecSimType_INT64:
         return sizeof(int64_t);
     }
     return 0;
 }
+
+size_t VecSimParams_GetDataSize(VecSimType type, size_t dim, VecSimMetric metric) {
+    size_t dataSize = VecSimType_sizeof(type) * dim;
+    if (type == VecSimType_INT8 && metric == VecSimMetric_Cosine) {
+        dataSize += sizeof(float); // For the norm
+    }
+    return dataSize;
+}

src/VecSim/utils/vec_utils.h

@@ -27,6 +27,7 @@ struct VecSimCommonStrings {
     static const char *FLOAT64_STRING;
     static const char *BFLOAT16_STRING;
     static const char *FLOAT16_STRING;
+    static const char *INT8_STRING;
     static const char *INT32_STRING;
     static const char *INT64_STRING;
 
@@ -90,3 +91,6 @@ const char *VecSimMetric_ToString(VecSimMetric vecsimMetric);
 const char *VecSimSearchMode_ToString(VecSearchMode vecsimSearchMode);
 
 size_t VecSimType_sizeof(VecSimType vecsimType);
+
+/** Returns the size in bytes of a stored or query vector */
+size_t VecSimParams_GetDataSize(VecSimType type, size_t dim, VecSimMetric metric);

src/VecSim/vec_sim.cpp

@@ -138,6 +138,9 @@ extern "C" void VecSim_Normalize(void *blob, size_t dim, VecSimType type) {
         spaces::GetNormalizeFunc<vecsim_types::bfloat16>()(blob, dim);
     } else if (type == VecSimType_FLOAT16) {
         spaces::GetNormalizeFunc<vecsim_types::float16>()(blob, dim);
+    } else if (type == VecSimType_INT8) {
+        // assuming blob is large enough to store the norm at the end of the vector
+        spaces::GetNormalizeFunc<int8_t>()(blob, dim);
     }
 }