Skip to content

Feature/rabitq submodule #465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
MaggieQi wants to merge 12 commits into from
Closed

Feature/rabitq submodule #465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
788f2cc
feat(SPANN): add official RaBitQ postings
lmccccc Jul 28, 2026
3459646
build: track official RaBitQ as submodule
lmccccc Jul 28, 2026
efa847b
refactor(SPANN): remove legacy global RaBitQ adapter
lmccccc Jul 31, 2026
f5b9235
refactor(SPANN): make RaBitQ a global quantizer
lmccccc Jul 31, 2026
4ece59a
clean some code
Jul 31, 2026
5ffc2e0
Feature/rabitq submodule (#454) (#456)
MaggieQi Jul 31, 2026
9e8180d
clean some code
Jul 31, 2026
18ebfce
Merge branch 'users/qiazh/pre-merge-tikv-bugfix' of https://github.co…
Jul 31, 2026
ba1699f
remove unused code
Jul 31, 2026
a1a7a5c
Add compact official RaBitQ integration (#459)
lmccccc Aug 21, 2026
fa9908a
Users/qiazh/pre merge tikv bugfix (#461)
MaggieQi Aug 21, 2026
8ad09af
Feature/rabitq submodule (#464)
lmccccc Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 98 additions & 83 deletions AnnService/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,21 +3,40 @@

set(AnnService ${PROJECT_SOURCE_DIR}/AnnService)
set(Zstd ${PROJECT_SOURCE_DIR}/ThirdParty/zstd)
set(RaBitQOfficial ${PROJECT_SOURCE_DIR}/ThirdParty/RaBitQOfficial)
set(RaBitQOfficialPatches ${PROJECT_SOURCE_DIR}/ThirdParty/RaBitQOfficialPatches)

if(NOT EXISTS ${RaBitQOfficial}/include/rabitqlib/defines.hpp)
message(FATAL_ERROR "RaBitQOfficial submodule is missing. Run: git submodule update --init --recursive")
endif()

include_directories(${AnnService})
include_directories(${Zstd}/lib)
include_directories(${RaBitQOfficialPatches}/include)
include_directories(${RaBitQOfficial}/include)

file(GLOB_RECURSE HDR_FILES ${AnnService}/inc/Core/*.h ${AnnService}/inc/Helper/*.h)
file(GLOB_RECURSE SRC_FILES ${AnnService}/src/Core/*.cpp ${AnnService}/src/Helper/*.cpp)

list(REMOVE_ITEM HDR_FILES
${AnnService}/inc/Core/Common/DistanceUtils.h
${AnnService}/inc/Core/Common/SIMDUtils.h
${AnnService}/inc/Core/Common/InstructionUtils.h
${AnnService}/inc/Core/Common/CommonUtils.h
)

list(REMOVE_ITEM SRC_FILES
${AnnService}/src/Core/Common/DistanceUtils.cpp
${AnnService}/src/Core/Common/SIMDUtils.cpp
${AnnService}/src/Core/Common/InstructionUtils.cpp
)

add_library (DistanceUtils STATIC
inc/Core/Common/DistanceUtils.h
inc/Core/Common/SIMDUtils.h
inc/Core/Common/InstructionUtils.h
inc/Core/Common/CommonUtils.h
src/Core/Common/DistanceUtils.cpp
src/Core/Common/SIMDUtils.cpp
src/Core/Common/InstructionUtils.cpp
)

if(${CMAKE_CXX_COMPILER_ID} STREQUAL "GNU")
target_compile_options(DistanceUtils PRIVATE -mavx2 -mavx -msse -msse2 -fPIC)
endif()

set(SPDK_LIBRARIES "")
if (SPDK)
set(Spdk ${PROJECT_SOURCE_DIR}/ThirdParty/spdk/build)
Expand Down Expand Up @@ -51,95 +70,91 @@ else()
)
endif()

list(REMOVE_ITEM HDR_FILES
${AnnService}/inc/Core/Common/DistanceUtils.h
${AnnService}/inc/Core/Common/SIMDUtils.h
${AnnService}/inc/Core/Common/InstructionUtils.h
${AnnService}/inc/Core/Common/CommonUtils.h
)

list(REMOVE_ITEM SRC_FILES
${AnnService}/src/Core/Common/DistanceUtils.cpp
${AnnService}/src/Core/Common/SIMDUtils.cpp
${AnnService}/src/Core/Common/InstructionUtils.cpp
)
set (RabitQ_LIBRARIES "")
if (RABITQ)
add_definitions(-DRABITQ)
set(RaBitQOfficial ${PROJECT_SOURCE_DIR}/ThirdParty/RaBitQOfficial)
set(RaBitQOfficialPatches ${PROJECT_SOURCE_DIR}/ThirdParty/RaBitQOfficialPatches)

add_library (DistanceUtils STATIC
inc/Core/Common/DistanceUtils.h
inc/Core/Common/SIMDUtils.h
inc/Core/Common/InstructionUtils.h
inc/Core/Common/CommonUtils.h
src/Core/Common/DistanceUtils.cpp
src/Core/Common/SIMDUtils.cpp
src/Core/Common/InstructionUtils.cpp
)
if(NOT EXISTS ${RaBitQOfficial}/include/rabitqlib/defines.hpp)
message(FATAL_ERROR "RaBitQOfficial submodule is missing. Run: git submodule update --init --recursive")
endif()

if(${CMAKE_CXX_COMPILER_ID} STREQUAL "GNU")
target_compile_options(DistanceUtils PRIVATE -mavx2 -mavx -msse -msse2 -fPIC)
endif()
include_directories(${RaBitQOfficialPatches}/include)
include_directories(${RaBitQOfficial}/include)

set(RABITQ_OFFICIAL_COMMON_SOURCES
${RaBitQOfficialPatches}/src/utils/cpu_features.cpp
${RaBitQOfficialPatches}/src/simd/dispatch.cpp
set(RABITQ_OFFICIAL_COMMON_SOURCES
${RaBitQOfficialPatches}/src/utils/cpu_features.cpp
${RaBitQOfficialPatches}/src/simd/dispatch.cpp
)
set(RABITQ_OFFICIAL_AVX2_SOURCES
${RaBitQOfficial}/src/simd/pack_excode_avx2.cpp
${RaBitQOfficial}/src/simd/space_excode_avx2.cpp
${RaBitQOfficial}/src/simd/space_avx2.cpp
${RaBitQOfficial}/src/simd/fastscan_avx2.cpp
${RaBitQOfficial}/src/simd/warmup_avx2.cpp
${RaBitQOfficial}/src/simd/rotator_avx2.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx2.cpp
set(RABITQ_OFFICIAL_AVX2_SOURCES
${RaBitQOfficial}/src/simd/pack_excode_avx2.cpp
${RaBitQOfficial}/src/simd/space_excode_avx2.cpp
${RaBitQOfficial}/src/simd/space_avx2.cpp
${RaBitQOfficial}/src/simd/fastscan_avx2.cpp
${RaBitQOfficial}/src/simd/warmup_avx2.cpp
${RaBitQOfficial}/src/simd/rotator_avx2.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx2.cpp
)
set(RABITQ_OFFICIAL_AVX512_SOURCES
${RaBitQOfficial}/src/simd/pack_excode_avx512.cpp
${RaBitQOfficial}/src/simd/space_excode_avx512.cpp
${RaBitQOfficial}/src/simd/space_avx512.cpp
${RaBitQOfficial}/src/simd/fastscan_avx512.cpp
${RaBitQOfficial}/src/simd/rotator_avx512.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx512_core.cpp
set(RABITQ_OFFICIAL_AVX512_SOURCES
${RaBitQOfficial}/src/simd/pack_excode_avx512.cpp
${RaBitQOfficial}/src/simd/space_excode_avx512.cpp
${RaBitQOfficial}/src/simd/space_avx512.cpp
${RaBitQOfficial}/src/simd/fastscan_avx512.cpp
${RaBitQOfficial}/src/simd/rotator_avx512.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx512_core.cpp
)
set(RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES
${RaBitQOfficial}/src/simd/warmup_avx512.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx512_popcnt.cpp
set(RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES
${RaBitQOfficial}/src/simd/warmup_avx512.cpp
${RaBitQOfficial}/src/index/hnsw_search_avx512_popcnt.cpp
)

set(RABITQ_OFFICIAL_SOURCES
${RABITQ_OFFICIAL_COMMON_SOURCES}
${RABITQ_OFFICIAL_AVX2_SOURCES}
set(RABITQ_OFFICIAL_SOURCES
${RABITQ_OFFICIAL_COMMON_SOURCES}
${RABITQ_OFFICIAL_AVX2_SOURCES}
)
if(NOT MSVC)
list(APPEND RABITQ_OFFICIAL_SOURCES
${RABITQ_OFFICIAL_AVX512_SOURCES}
${RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES})
endif()
if(NOT MSVC)
list(APPEND RABITQ_OFFICIAL_SOURCES
${RABITQ_OFFICIAL_AVX512_SOURCES}
${RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES})
endif()

add_library(RaBitQOfficialCore STATIC ${RABITQ_OFFICIAL_SOURCES})
target_include_directories(RaBitQOfficialCore PUBLIC
${RaBitQOfficialPatches}/include
${RaBitQOfficial}/include)
set_target_properties(RaBitQOfficialCore PROPERTIES POSITION_INDEPENDENT_CODE ON)
if(${CMAKE_CXX_COMPILER_ID} STREQUAL "GNU")
target_compile_options(RaBitQOfficialCore PRIVATE -fopenmp)
set_source_files_properties(${RABITQ_OFFICIAL_AVX2_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
set_source_files_properties(${RABITQ_OFFICIAL_AVX512_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512bw -mavx512dq -mfma")
set_source_files_properties(${RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512bw -mavx512dq -mavx512vpopcntdq -mfma")
elseif(MSVC)
target_compile_definitions(RaBitQOfficialCore PRIVATE RABITQ_NO_AVX512=1)
target_compile_options(RaBitQOfficialCore PRIVATE
/openmp
"/FI${RaBitQOfficialPatches}/include/rabitqlib/utils/compiler_compat.hpp")
set_source_files_properties(${RABITQ_OFFICIAL_AVX2_SOURCES}
PROPERTIES COMPILE_FLAGS "/arch:AVX2")
add_library(RaBitQOfficialCore STATIC ${RABITQ_OFFICIAL_SOURCES})
target_include_directories(RaBitQOfficialCore PUBLIC
${RaBitQOfficialPatches}/include
${RaBitQOfficial}/include)
set_target_properties(RaBitQOfficialCore PROPERTIES POSITION_INDEPENDENT_CODE ON)
if(${CMAKE_CXX_COMPILER_ID} STREQUAL "GNU")
target_compile_options(RaBitQOfficialCore PRIVATE -fopenmp)
set_source_files_properties(${RABITQ_OFFICIAL_AVX2_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
set_source_files_properties(${RABITQ_OFFICIAL_AVX512_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512bw -mavx512dq -mfma")
set_source_files_properties(${RABITQ_OFFICIAL_AVX512_POPCNT_SOURCES}
PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512bw -mavx512dq -mavx512vpopcntdq -mfma")
elseif(MSVC)
target_compile_definitions(RaBitQOfficialCore PRIVATE RABITQ_NO_AVX512=1)
target_compile_options(RaBitQOfficialCore PRIVATE
/openmp
"/FI${RaBitQOfficialPatches}/include/rabitqlib/utils/compiler_compat.hpp")
set_source_files_properties(${RABITQ_OFFICIAL_AVX2_SOURCES}
PROPERTIES COMPILE_FLAGS "/arch:AVX2")
endif()
set (RabitQ_LIBRARIES RabitQOfficialCore)
else()
list(REMOVE_ITEM HDR_FILES
${AnnService}/inc/Core/Common/RaBitQQuantizer.h
)

list(REMOVE_ITEM SRC_FILES
${AnnService}/src/Core/Common/RaBitQQuantizer.cpp
)
endif()

add_library (SPTAGLib SHARED ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES})
target_link_libraries (SPTAGLib DistanceUtils RaBitQOfficialCore ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
target_link_libraries (SPTAGLib DistanceUtils ${RabitQ_LIBRARIES} ${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 RaBitQOfficialCore ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})
target_link_libraries (SPTAGLibStatic DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES})

if (MSVC)
# SPANNIndex.cpp can exceed COFF section limits in Debug without /bigobj.
Expand Down
3 changes: 2 additions & 1 deletion AnnService/inc/Core/BKT/ParameterDefinitionList.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ DefineBKTParameter(m_pTrees.m_iBKTKmeansK, int, 32L, "BKTKmeansK")
DefineBKTParameter(m_pTrees.m_iBKTLeafSize, int, 8L, "BKTLeafSize")
DefineBKTParameter(m_pTrees.m_iSamples, int, 1000L, "Samples")
DefineBKTParameter(m_pTrees.m_fBalanceFactor, float, 100.0F, "BKTLambdaFactor")
DefineBKTParameter(m_pTrees.m_parallelBuild, bool, false, "ParallelBKTBuild")
DefineBKTParameter(m_pTrees.m_parallelBuild, bool, true, "ParallelBKTBuild")

DefineBKTParameter(m_pGraph.m_iTPTNumber, int, 32L, "TPTNumber")
DefineBKTParameter(m_pGraph.m_iTPTLeafSize, int, 2000L, "TPTLeafSize")
Expand All @@ -37,6 +37,7 @@ DefineBKTParameter(m_pGraph.m_iGPURefineDepth, int, 30, "GPURefineDepth") // Dep
DefineBKTParameter(m_pGraph.m_iGPULeafSize, int, 500, "GPULeafSize")
DefineBKTParameter(m_pGraph.m_iheadNumGPUs, int, 1, "HeadNumGPUs")
DefineBKTParameter(m_pGraph.m_iTPTBalanceFactor, int, 2, "TPTBalanceFactor")
DefineBKTParameter(m_pGraph.m_TPTParallelBuild, bool, true, "ParallelTPTBuild")

DefineBKTParameter(m_iNumberOfThreads, int, 1L, "NumberOfThreads")
DefineBKTParameter(m_iDistCalcMethod, SPTAG::DistCalcMethod, SPTAG::DistCalcMethod::Cosine, "DistCalcMethod")
Expand Down
29 changes: 14 additions & 15 deletions AnnService/inc/Core/Common/FineGrainedLock.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,29 +22,29 @@ namespace SPTAG
~FineGrainedLock() {}

std::mutex& operator[](SizeType idx) {
unsigned index = hash_func((unsigned)idx);
std::uint64_t index = hash_func((std::uint64_t)idx);
return m_locks[index];
}

const std::mutex& operator[](SizeType idx) const {
unsigned index = hash_func((unsigned)idx);
std::uint64_t index = hash_func((std::uint64_t)idx);
return m_locks[index];
}

static inline unsigned hash_func(unsigned idx)
static inline std::uint64_t hash_func(std::uint64_t idx)
{
return ((unsigned)(idx * 99991) + _rotl(idx, 2) + 101) & PoolSize;
return (idx * 99991 + _rotl64(idx, 2) + 101) & PoolSize;
}

private:
static const int PoolSize = 32767;
static const std::uint64_t PoolSize = 32767;
std::unique_ptr<std::mutex[]> m_locks;
};

class FineGrainedRWLock {
public:
FineGrainedRWLock() {
m_buckets.reset(new Bucket[BucketCount]);
m_buckets.reset(new Bucket[BucketSize + 1]);
}
~FineGrainedRWLock() {}

Expand All @@ -56,10 +56,10 @@ namespace SPTAG
return GetLock(idx);
}

static inline unsigned hash_func(unsigned idx)
{
static inline SizeType hash_func(SizeType idx)
{
return idx;
}
}
private:
struct Bucket {
std::mutex mutex;
Expand All @@ -76,17 +76,16 @@ namespace SPTAG
return *iter->second;
}

static inline unsigned BucketIndex(SizeType idx)
static inline std::uint64_t BucketIndex(SizeType idx)
{
unsigned key = static_cast<unsigned>(idx);
return ((unsigned)(key * 99991) + _rotl(key, 2) + 101) & BucketMask;
std::uint64_t key = (std::uint64_t)idx;
return (key * 99991 + _rotl64(key, 2) + 101) & BucketSize;
}

static const int BucketMask = 32767;
static const int BucketCount = BucketMask + 1;
static const std::uint64_t BucketSize = 32767;
mutable std::unique_ptr<Bucket[]> m_buckets;
};
}
}

#endif // _SPTAG_COMMON_FINEGRAINEDLOCK_H_
#endif // _SPTAG_COMMON_FINEGRAINEDLOCK_H_
Loading