Feature/rabitq auto tuning

AnnService/CMakeLists.txt

@@ -144,17 +144,23 @@ if (RABITQ)
 else()
     list(REMOVE_ITEM HDR_FILES
 	    ${AnnService}/inc/Core/Common/RaBitQQuantizer.h
+	    ${AnnService}/inc/Core/Common/RaBitQAutoTuner.h
     )
 
     list(REMOVE_ITEM SRC_FILES
 	    ${AnnService}/src/Core/Common/RaBitQQuantizer.cpp
+	    ${AnnService}/src/Core/Common/RaBitQAutoTuner.cpp
     )
 endif()
 
 add_library (SPTAGLib SHARED ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES})
-target_link_libraries (SPTAGLib DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
+target_link_libraries (SPTAGLib DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
 add_library (SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES})
-target_link_libraries (SPTAGLibStatic DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
+target_link_libraries (SPTAGLibStatic DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
+if (RABITQ)
+    target_link_libraries(SPTAGLib RaBitQOfficialCore)
+    target_link_libraries(SPTAGLibStatic RaBitQOfficialCore)
+endif()
 
 if (MSVC)
     # SPANNIndex.cpp can exceed COFF section limits in Debug without /bigobj.

AnnService/inc/Core/Common/IQuantizer.h

@@ -60,6 +60,8 @@ namespace SPTAG
 
             virtual float* GetL2DistanceTables() = 0;
 
+            virtual bool QuantizeForIndexBuild() const { return true; }
+
             template<typename T>
             T* GetCodebooks();
         };

AnnService/inc/Core/Common/OPQQuantizer.h

@@ -53,6 +53,10 @@ namespace SPTAG
                 return GetEnumValueType<T>();
             }
 
+            bool QuantizeForIndexBuild() const override
+            {
+                return false;
+            }
 
         protected:
             using PQQuantizer<OPQMatrixType>::m_NumSubvectors;
@@ -200,7 +204,11 @@ namespace SPTAG
         inline void OPQQuantizer<T>::m_VectorMatrixMultiply(OPQMatrixType* mat, const OPQMatrixType* vec, O* mat_vec) const
         {
             for (int i = 0; i < m_matrixDim; i++) {
-                mat_vec[i] = (O)(m_base - m_fdot(vec, mat, m_matrixDim));
+                OPQMatrixType value = 0;
+                for (int j = 0; j < m_matrixDim; ++j) {
+                    value += vec[j] * mat[j];
+                }
+                mat_vec[i] = static_cast<O>(value);
                 mat += m_matrixDim;
             }
         }

AnnService/inc/Core/Common/RaBitQAutoTuner.h

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "inc/Core/Common/RaBitQQuantizer.h"
+#include "inc/Helper/SimpleIniReader.h"
+#include "inc/Helper/VectorSetReader.h"
+
+#include <functional>
+#include <string>
+#include <vector>
+
+namespace SPTAG
+{
+namespace COMMON
+{
+
+struct RaBitQAutoTuneResult
+{
+    int selectedBits = 0;
+    float recall = 0.0F;
+    SizeType vectorCount = 0;
+    DimensionType codeDimension = 0;
+    std::string quantizerPath;
+    std::string vectorPath;
+    std::shared_ptr<RaBitQQuantizer> quantizer;
+};
+
+class RaBitQAutoTuner
+{
+public:
+    using BitEvaluator = std::function<ErrorCode(int, float&)>;
+
+    static bool IsEnabled(const Helper::IniReader& p_config);
+    static ErrorCode Run(Helper::IniReader& p_config,
+                         const std::string& p_outputFolder,
+                         RaBitQAutoTuneResult& p_result,
+                         std::string& p_error);
+
+    static ErrorCode SelectMinimumBits(float p_targetRecall,
+                                       const BitEvaluator& p_evaluator,
+                                       int& p_selectedBits,
+                                       float& p_selectedRecall);
+    static ErrorCode ValidateTruth(const std::vector<std::vector<SizeType>>& p_truth,
+                                   SizeType p_baseCount,
+                                   SizeType p_queryCount,
+                                   int p_resultCount,
+                                   std::string& p_error);
+    static float RecallAtK(const std::vector<SizeType>& p_exact,
+                           const std::vector<SizeType>& p_reranked,
+                           int p_resultCount);
+};
+
+} // namespace COMMON
+} // namespace SPTAG

AnnService/inc/Core/Common/RaBitQQuantizer.h

@@ -25,6 +25,10 @@ class RaBitQQuantizer : public IQuantizer
     RaBitQQuantizer(DimensionType p_dimension, int p_bits, bool p_normalize);
 
     ErrorCode Train(const std::shared_ptr<VectorSet>& p_vectors);
+    ErrorCode BeginTraining();
+    ErrorCode AddTrainingBatch(const std::shared_ptr<VectorSet>& p_vectors);
+    ErrorCode FinishTraining();
+    std::shared_ptr<RaBitQQuantizer> CreateWithBits(int p_bits) const;
 
     float L2Distance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override;
     float CosineDistance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override;
@@ -44,10 +48,12 @@ class RaBitQQuantizer : public IQuantizer
     DimensionType GetNumSubvectors() const override;
     int GetBase() const override;
     float* GetL2DistanceTables() override;
+    bool QuantizeForIndexBuild() const override { return false; }
 
     DimensionType Dimension() const { return m_dimension; }
     int Bits() const { return m_bits; }
     bool Ready() const;
+    bool Trained() const { return m_trained; }
 
 private:
     struct ModelHeader
@@ -90,6 +96,9 @@ class RaBitQQuantizer : public IQuantizer
     rabitqlib::quant::RabitqConfig m_quantizer_config;
     rabitqlib::ex_ipfunc m_ip_func = nullptr;
     std::vector<float> m_centroid;
+    std::vector<double> m_training_sum;
+    std::uint64_t m_training_count = 0;
+    bool m_trained = false;
 };
 
 } // namespace COMMON

AnnService/inc/Core/Common/TruthSet.h

@@ -169,6 +169,11 @@ namespace SPTAG
                 float meanrecall = 0, minrecall = MaxDist, maxrecall = 0, stdrecall = 0, meanmrr = 0;
                 std::vector<float> thisrecall(NumQuerys, 0);
                 std::unique_ptr<bool[]> visited(new bool[K]);
+                const bool compareDistanceTies =
+                    querySet != nullptr &&
+                    vectorSet != nullptr &&
+                    querySet->GetValueType() == GetEnumValueType<T>() &&
+                    vectorSet->GetValueType() == GetEnumValueType<T>();
                 for (SizeType i = 0; i < NumQuerys; i++)
                 {
                     int minpos = K;
@@ -186,7 +191,7 @@ namespace SPTAG
                                 if (j < minpos) minpos = j;
                                 break;
                             }
-                            else if (vectorSet != nullptr) {
+                            else if (compareDistanceTies) {
                                 float dist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(results[i].GetResult(j)->VID), vectorSet->Dimension(), index->GetDistCalcMethod());
                                 float truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), vectorSet->Dimension(), index->GetDistCalcMethod());
                                 if (index->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine && fabs(dist - truthDist) < Epsilon) {
@@ -213,7 +218,7 @@ namespace SPTAG
                         std::vector<NodeDistPair> truthvec;
                         for (SizeType id : truth[i]) {
                             float truthDist = 0.0;
-                            if (vectorSet != nullptr) {
+                            if (compareDistanceTies) {
                                 truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), querySet->Dimension(), index->GetDistCalcMethod());
                             }
                             truthvec.emplace_back(id, truthDist);

AnnService/inc/Core/Common/WorkSpace.h

@@ -146,13 +146,21 @@ namespace SPTAG
 
             inline void DoubleSize()
             {
+                const std::uint64_t oldPoolSize = m_poolSize;
+                const bool hadSecondHash = m_secondHash;
                 std::uint64_t new_poolSize = ((m_poolSize + 1) << 1) - 1; 
                 SizeType* new_hashTable = new SizeType[(new_poolSize + 1) * 2];
                 memset(new_hashTable, 0, sizeof(SizeType) * (new_poolSize + 1) * 2);
 
                 m_secondHash = false;
-                for (std::uint64_t i = 0; i <= new_poolSize; i++)
+                for (std::uint64_t i = 0; i <= oldPoolSize; i++)
                     if (m_hashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, m_hashTable[i]);
+                if (hadSecondHash)
+                {
+                    SizeType* secondHashTable = m_hashTable.get() + oldPoolSize + 1;
+                    for (std::uint64_t i = 0; i <= oldPoolSize; i++)
+                        if (secondHashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, secondHashTable[i]);
+                }
 
                 m_exp++;
                 m_poolSize = new_poolSize;