[MOD-14957] Add SQ8 support for tiered HNSW

src/VecSim/algorithms/brute_force/brute_force_single.h

@@ -155,7 +155,8 @@ int BruteForceIndex_Single<DataType, DistType>::addVector(const void *vector_dat
     // Check if label already exists, so it is an update operation.
     if (optionalID != this->labelToIdLookup.end()) {
         idType id = optionalID->second;
-        this->vectors->updateElement(id, vector_data);
+        auto processed_blob = this->preprocessForStorage(vector_data);
+        this->vectors->updateElement(id, processed_blob.get());
         return 0;
     }
 

src/VecSim/algorithms/hnsw/hnsw.h

@@ -251,6 +251,7 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
     void unlockIndexDataGuard() const;
     void lockSharedIndexDataGuard() const;
     void unlockSharedIndexDataGuard() const;
+    std::shared_lock<std::shared_mutex> acquireSharedIndexDataGuard() const;
     void lockNodeLinks(idType node_id) const;
     void unlockNodeLinks(idType node_id) const;
     VisitedNodesHandler *getVisitedList() const;
@@ -545,6 +546,12 @@ void HNSWIndex<DataType, DistType>::unlockSharedIndexDataGuard() const {
     indexDataGuard.unlock_shared();
 }
 
+template <typename DataType, typename DistType>
+std::shared_lock<std::shared_mutex>
+HNSWIndex<DataType, DistType>::acquireSharedIndexDataGuard() const {
+    return std::shared_lock<std::shared_mutex>(indexDataGuard);
+}
+
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::lockNodeLinks(idType node_id) const {
     elementLocks[node_id].lock();

src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h

@@ -68,6 +68,9 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarked
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test)
 
+INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestSQ8)
+INDEX_TEST_FRIEND_CLASS(SQ8TieredHNSWTest)
+
 friend class CommonAPITest_SearchDifferentScores_Test;
 friend class BF16TieredTest;
 friend class FP16TieredTest;

src/VecSim/index_factories/components/components_factory.h

@@ -31,6 +31,58 @@ CreateIndexComponents(std::shared_ptr<VecSimAllocator> allocator, VecSimMetric m
     return {indexCalculator, preprocessors};
 }
 
+// Asymmetric dispatch reports alignment for the stored operand only. Ask the query type's
+// dispatcher for the query allocation alignment.
+template <typename DataType>
+[[nodiscard]] unsigned char GetQueryAlignment(VecSimMetric metric, size_t dim) {
+    unsigned char alignment = 0;
+    spaces::GetDistFunc<DataType, float>(metric, dim, &alignment);
+    return alignment;
+}
+
+template <typename DataType, VecSimMetric Metric>
+IndexComponents<DataType, float>
+CreateSQ8IndexComponents(const std::shared_ptr<VecSimAllocator> &allocator, size_t dim,
+                         const float *mean_ptr) {
+    const bool with_norm = mean_ptr != nullptr;
+    unsigned char storage_alignment = 0, asym_storage_alignment = 0;
+
+    // Graph construction compares two stored SQ8 blobs; search compares a stored blob with a
+    // DataType query. Both dispatchers report alignment for the stored operand.
+    auto sym_func = spaces::GetDistFunc<vecsim_types::sq8, float>(Metric, dim, &storage_alignment);
+    auto asym_func = spaces::GetDistFunc<vecsim_types::sq8, float, DataType>(
+        Metric, dim, &asym_storage_alignment);
+    storage_alignment = spaces::combineAlignments(storage_alignment, asym_storage_alignment);
+    const unsigned char query_alignment = GetQueryAlignment<DataType>(Metric, dim);
+
+    PreprocessorInterface *pp = nullptr;
+    IndexCalculatorInterface<float> *calc = nullptr;
+
+    if (with_norm) {
+        vecsim_stl::vector<float> mean_vec(allocator);
+        mean_vec.assign(mean_ptr, mean_ptr + dim);
+
+        float mean_sum_squares = 0.0f;
+        for (float v : mean_vec) {
+            mean_sum_squares += v * v;
+        }
+
+        pp = new (allocator) QuantPreprocessor<DataType, Metric, true>(allocator, dim, mean_vec);
+        calc = new (allocator) DistanceCalculatorWithNorm<DataType, float, Metric>(
+            allocator, asym_func, sym_func, mean_sum_squares);
+    } else {
+        pp = new (allocator) QuantPreprocessor<DataType, Metric>(allocator, dim);
+        calc = new (allocator) DistanceCalculatorCommon<float>(allocator, sym_func, asym_func);
+    }
+
+    auto *container = new (allocator)
+        MultiPreprocessorsContainer<DataType, 1>(allocator, query_alignment, storage_alignment);
+    [[maybe_unused]] const int ret = container->addPreprocessor(pp);
+    assert(ret != -1 && "SQ8 preprocessor was not added correctly");
+
+    return {calc, container};
+}
+
 template <typename DataType, typename DistType>
 size_t EstimateComponentsMemory(VecSimMetric metric, bool is_normalized) {
     size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize();

src/VecSim/index_factories/hnsw_factory.cpp

@@ -43,15 +43,6 @@ template <VecSimMetric Metric>
                      : sq8::storage_bytes_count<Metric, false>(dim);
 }
 
-// Asymmetric dispatch reports alignment for the stored operand only. Ask the query type's
-// dispatcher for the query allocation alignment.
-template <typename DataType>
-[[nodiscard]] unsigned char GetQueryAlignment(VecSimMetric metric, size_t dim) {
-    unsigned char alignment = 0;
-    spaces::GetDistFunc<DataType, float>(metric, dim, &alignment);
-    return alignment;
-}
-
 // Cosine over pre-normalized vectors is computed as inner product.
 [[nodiscard]] constexpr VecSimMetric ResolveSQ8Metric(VecSimMetric metric, bool is_normalized) {
     return (is_normalized && metric == VecSimMetric_Cosine) ? VecSimMetric_IP : metric;
@@ -79,48 +70,13 @@ template <typename DataType>
 template <typename DataType, VecSimMetric Metric>
 VecSimIndex *NewIndex_SQ8(const HNSWParams *hnswParams, AbstractIndexInitParams abstractInitParams,
                           const float *mean_ptr) {
-    auto &allocator = abstractInitParams.allocator;
-    const size_t dim = abstractInitParams.dim;
     const bool with_norm = mean_ptr != nullptr;
-    unsigned char storage_alignment = 0, asym_storage_alignment = 0;
-
-    abstractInitParams.storedDataSize = GetSQ8StoredDataSize<Metric>(dim, with_norm);
+    abstractInitParams.storedDataSize =
+        GetSQ8StoredDataSize<Metric>(abstractInitParams.dim, with_norm);
     abstractInitParams.isQuantized = true;
 
-    // Graph construction compares two stored SQ8 blobs; search compares a stored blob with a
-    // DataType query. Both dispatchers report alignment for the stored operand.
-    auto sym_func = spaces::GetDistFunc<sq8, float>(Metric, dim, &storage_alignment);
-    auto asym_func =
-        spaces::GetDistFunc<sq8, float, DataType>(Metric, dim, &asym_storage_alignment);
-    storage_alignment = spaces::combineAlignments(storage_alignment, asym_storage_alignment);
-    const unsigned char query_alignment = GetQueryAlignment<DataType>(Metric, dim);
-
-    PreprocessorInterface *pp = nullptr;
-    IndexCalculatorInterface<float> *calc = nullptr;
-
-    if (with_norm) {
-        vecsim_stl::vector<float> mean_vec(allocator);
-        mean_vec.assign(mean_ptr, mean_ptr + dim);
-
-        float mean_sum_squares = 0.0f;
-        for (float v : mean_vec) {
-            mean_sum_squares += v * v;
-        }
-
-        pp = new (allocator) QuantPreprocessor<DataType, Metric, true>(allocator, dim, mean_vec);
-        calc = new (allocator) DistanceCalculatorWithNorm<DataType, float, Metric>(
-            allocator, asym_func, sym_func, mean_sum_squares);
-    } else {
-        pp = new (allocator) QuantPreprocessor<DataType, Metric>(allocator, dim);
-        calc = new (allocator) DistanceCalculatorCommon<float>(allocator, sym_func, asym_func);
-    }
-
-    auto *container = new (allocator)
-        MultiPreprocessorsContainer<DataType, 1>(allocator, query_alignment, storage_alignment);
-    [[maybe_unused]] const int ret = container->addPreprocessor(pp);
-    assert(ret != -1 && "SQ8 preprocessor was not added correctly");
-
-    IndexComponents<DataType, float> components{calc, container};
+    IndexComponents<DataType, float> components = CreateSQ8IndexComponents<DataType, Metric>(
+        abstractInitParams.allocator, abstractInitParams.dim, mean_ptr);
     return NewIndex_ChooseMultiOrSingle<DataType, float>(hnswParams, abstractInitParams,
                                                          components);
 }

src/VecSim/index_factories/tiered_factory.cpp

@@ -40,21 +40,57 @@ static inline BFParams NewBFParams(const TieredIndexParams *params) {
 template <typename DataType, typename DistType = DataType>
 inline VecSimIndex *NewIndex(const TieredIndexParams *params) {
 
-    // initialize hnsw index
-    // Normalization is done by the frontend index.
-    auto *hnsw_index = reinterpret_cast<HNSWIndex<DataType, DistType> *>(
-        HNSWFactory::NewIndex(params->primaryIndexParams, true));
-    // initialize brute force index
+    const auto &hnsw_params = params->primaryIndexParams->algoParams.hnswParams;
+    bool defer_backend = false;
+
+    if (hnsw_params.quantType != VecSimQuant_NONE) {
+        constexpr bool supports_quantization =
+            std::is_same_v<DistType, float> &&
+            (std::is_same_v<DataType, float> || std::is_same_v<DataType, float16>);
+
+        if (!supports_quantization || hnsw_params.quantType != VecSimQuant_SQ8) {
+            return nullptr;
+        }
+
+        const bool with_norm =
+            params->specificParams.tieredHnswParams.QuantNormalizationSetSize > 0;
+        const bool supports_with_norm =
+            !(std::is_same_v<DataType, float16> && hnsw_params.metric == VecSimMetric_L2);
+
+        if (with_norm) {
+            if (!supports_with_norm) {
+                return nullptr;
+            } else {
+                defer_backend = true;
+            }
+        }
+    }
+
+    HNSWIndex<DataType, DistType> *hnsw_index = nullptr;
+    if (!defer_backend) {
+        // Normalization is done by the frontend index.
+        hnsw_index = reinterpret_cast<HNSWIndex<DataType, DistType> *>(
+            HNSWFactory::NewIndex(params->primaryIndexParams, true));
+    }
 
     BFParams bf_params = NewBFParams(params);
 
     AbstractIndexInitParams abstractInitParams =
         VecSimFactory::NewAbstractInitParams(&bf_params, params->primaryIndexParams->logCtx, false);
-    assert(hnsw_index->getInputBlobSize() == abstractInitParams.storedDataSize);
-    assert(hnsw_index->getStoredDataSize() == abstractInitParams.storedDataSize);
+    if (hnsw_index) {
+        assert(hnsw_index->getInputBlobSize() == abstractInitParams.storedDataSize);
+        if (hnsw_params.quantType == VecSimQuant_NONE) {
+            assert(hnsw_index->getStoredDataSize() == abstractInitParams.storedDataSize);
+        }
+    }
     auto frontendIndex = static_cast<BruteForceIndex<DataType, DistType> *>(
         BruteForceFactory::NewIndex(&bf_params, abstractInitParams, false));
 
+    if (hnsw_params.quantType == VecSimQuant_SQ8 && hnsw_params.dim < 64) {
+        frontendIndex->log(VecSimCommonStrings::LOG_WARNING_STRING,
+                           "SQ8 compression is not recommended for dimensions below 64");
+    }
+
     // Create new tiered hnsw index
     std::shared_ptr<VecSimAllocator> management_layer_allocator =
         VecSimAllocator::newVecsimAllocator();
@@ -66,17 +102,33 @@ inline VecSimIndex *NewIndex(const TieredIndexParams *params) {
 inline size_t EstimateInitialSize(const TieredIndexParams *params) {
     HNSWParams hnsw_params = params->primaryIndexParams->algoParams.hnswParams;
 
-    // Keep size estimation consistent with NewIndex, which rejects quantized tiered indexes.
-    if (hnsw_params.quantType != VecSimQuant_NONE) {
-        throw std::invalid_argument("Quantization is not supported for tiered HNSW indexes");
+    size_t est = 0;
+
+    const bool defer_backend =
+        hnsw_params.quantType != VecSimQuant_NONE &&
+        params->specificParams.tieredHnswParams.QuantNormalizationSetSize > 0;
+
+    if (defer_backend) {
+        // Set quantParams non-null to indicate HNSW SQ8 with_norm index
+        static char dummy;
+        hnsw_params.quantParams = &dummy;
     }
 
-    // Add size estimation of VecSimTieredIndex sub indexes.
-    // Normalization is done by the frontend index.
-    size_t est = HNSWFactory::EstimateInitialSize(&hnsw_params, true);
+    // HNSWFactory::EstimateInitialSize will throw if the parameters are invalid
+    size_t est_backend = HNSWFactory::EstimateInitialSize(&hnsw_params, true);
 
-    // Management layer allocator overhead.
     size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize();
+
+    if (defer_backend) {
+        // Add size of SQ accumulation buffer
+        est += allocations_overhead + hnsw_params.dim * sizeof(double);
+    } else {
+        // Add size estimation of VecSimTieredIndex sub indexes.
+        // Normalization is done by the frontend index.
+        est += est_backend;
+    }
+
+    // Management layer allocator overhead.
     est += sizeof(VecSimAllocator) + allocations_overhead;
 
     // Size of the TieredHNSWIndex struct.
@@ -100,12 +152,6 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) {
 }
 
 VecSimIndex *NewIndex(const TieredIndexParams *params) {
-    // The brute-force frontend is not quantized, so an SQ8 primary index would use an incompatible
-    // stored-vector layout.
-    if (params->primaryIndexParams->algoParams.hnswParams.quantType != VecSimQuant_NONE) {
-        return nullptr;
-    }
-
     // Tiered index that contains HNSW index as primary index
     VecSimType type = params->primaryIndexParams->algoParams.hnswParams.type;
     if (type == VecSimType_FLOAT32) {
@@ -247,7 +293,15 @@ size_t EstimateElementSize(const TieredIndexParams *params) {
     // Match HNSW's element estimator, which leaves validation to NewIndex.
     size_t est = 0;
     if (params->primaryIndexParams->algo == VecSimAlgo_HNSWLIB) {
-        est = HNSWFactory::EstimateElementSize(&params->primaryIndexParams->algoParams.hnswParams);
+        HNSWParams hnsw_params = params->primaryIndexParams->algoParams.hnswParams;
+        if (hnsw_params.quantType != VecSimQuant_NONE &&
+            params->specificParams.tieredHnswParams.QuantNormalizationSetSize > 0) {
+
+            // Set quantParams non-null to indicate HNSW SQ8 with_norm index
+            static char dummy;
+            hnsw_params.quantParams = &dummy;
+        }
+        est = HNSWFactory::EstimateElementSize(&hnsw_params);
     }
     if (params->primaryIndexParams->algo == VecSimAlgo_SVS) {
         est = SVSFactory::EstimateElementSize(&params->primaryIndexParams->algoParams.svsParams);

src/VecSim/vec_sim_common.h

@@ -229,6 +229,9 @@ typedef struct {
 typedef struct {
     size_t swapJobThreshold; // The minimum number of swap jobs to accumulate before applying
                              // all the ready swap jobs in a batch.
+    size_t QuantNormalizationSetSize; // Number of vectors to accumulate before SQ initialization.
+                                      // 0 = skip accumulation phase (naive SQ8, no mean).
+                                      // Max: 100 * DEFAULT_BLOCK_SIZE (102400).
 } TieredHNSWParams;
 
 // A struct that contains HNSW Disk tiered index specific params.