Skip to content

Commit 2e16ede

Browse files
introduce RVV platform and its ScalarQuantizer
Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>
1 parent e23c661 commit 2e16ede

9 files changed

Lines changed: 493 additions & 0 deletions

File tree

faiss/CMakeLists.txt

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,16 @@ set(FAISS_SIMD_SVE_SRC
4141
impl/pq_code_distance/pq_code_distance-sve.cpp
4242
utils/simd_impl/distances_arm_sve.cpp
4343
)
44+
set(FAISS_SIMD_RVV_SRC
45+
impl/scalar_quantizer/sq-rvv.cpp
46+
)
4447
# Select SIMD sources based on target architecture
4548
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)")
4649
set(FAISS_SIMD_SRC ${FAISS_SIMD_AVX2_SRC} ${FAISS_SIMD_AVX512_SRC})
4750
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
4851
set(FAISS_SIMD_SRC ${FAISS_SIMD_NEON_SRC} ${FAISS_SIMD_SVE_SRC})
52+
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)")
53+
set(FAISS_SIMD_SRC ${FAISS_SIMD_RVV_SRC})
4954
else()
5055
set(FAISS_SIMD_SRC "")
5156
endif()
@@ -477,6 +482,14 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
477482
TARGET_DIRECTORY faiss
478483
PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+sve"
479484
)
485+
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)")
486+
target_compile_definitions(faiss PRIVATE COMPILE_SIMD_RISCV_RVV)
487+
if(FAISS_SIMD_RVV_SRC)
488+
set_source_files_properties(${FAISS_SIMD_RVV_SRC}
489+
TARGET_DIRECTORY faiss
490+
PROPERTIES COMPILE_OPTIONS "-march=rv64gcv_zvfhmin;-mabi=lp64d"
491+
)
492+
endif()
480493
endif()
481494
endif()
482495
endif()
@@ -489,6 +502,19 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
489502
target_sources(faiss PRIVATE ${FAISS_SIMD_NEON_SRC})
490503
endif()
491504

505+
# RVV is the baseline SIMD on rv64 builds compiled with rv64gcv. Compile RVV
506+
# sources into the main faiss target, mirroring the ARM NEON story on aarch64.
507+
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(riscv64|riscv)")
508+
target_compile_definitions(faiss PRIVATE COMPILE_SIMD_RISCV_RVV)
509+
target_sources(faiss PRIVATE ${FAISS_SIMD_RVV_SRC})
510+
if(NOT WIN32 AND FAISS_SIMD_RVV_SRC)
511+
set_source_files_properties(${FAISS_SIMD_RVV_SRC}
512+
TARGET_DIRECTORY faiss
513+
PROPERTIES COMPILE_OPTIONS "-march=rv64gcv_zvfhmin;-mabi=lp64d"
514+
)
515+
endif()
516+
endif()
517+
492518
if(FAISS_ENABLE_SVS)
493519
find_package(svs_runtime REQUIRED)
494520

faiss/impl/ScalarQuantizer.h

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,25 @@ struct ScalarQuantizer : Quantizer {
102102

103103
virtual float query_to_code(const uint8_t* code) const = 0;
104104

105+
/// Compute four query-to-code distances in one call. Default loops
106+
/// query_to_code four times; per-SIMD specializations may batch the
107+
/// inner dim loop across the four codes to amortize query state and
108+
/// expose ILP across independent accumulators.
109+
virtual void query_to_codes_batch_4(
110+
const uint8_t* code_0,
111+
const uint8_t* code_1,
112+
const uint8_t* code_2,
113+
const uint8_t* code_3,
114+
float& dis0,
115+
float& dis1,
116+
float& dis2,
117+
float& dis3) const {
118+
dis0 = query_to_code(code_0);
119+
dis1 = query_to_code(code_1);
120+
dis2 = query_to_code(code_2);
121+
dis3 = query_to_code(code_3);
122+
}
123+
105124
float distance_to_code(const uint8_t* code) final {
106125
return query_to_code(code);
107126
}

faiss/impl/scalar_quantizer/sq-avx2.cpp

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,43 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
387387
float query_to_code(const uint8_t* code) const final {
388388
return compute_distance(q, code);
389389
}
390+
391+
void query_to_codes_batch_4(
392+
const uint8_t* code_0,
393+
const uint8_t* code_1,
394+
const uint8_t* code_2,
395+
const uint8_t* code_3,
396+
float& dis0,
397+
float& dis1,
398+
float& dis2,
399+
float& dis3) const final {
400+
Similarity sim0(q);
401+
Similarity sim1(q);
402+
Similarity sim2(q);
403+
Similarity sim3(q);
404+
405+
sim0.begin_8();
406+
sim1.begin_8();
407+
sim2.begin_8();
408+
sim3.begin_8();
409+
410+
for (size_t i = 0; i < quant.d; i += 8) {
411+
const int ii = static_cast<int>(i);
412+
simd8float32 xi0 = quant.reconstruct_8_components(code_0, ii);
413+
simd8float32 xi1 = quant.reconstruct_8_components(code_1, ii);
414+
simd8float32 xi2 = quant.reconstruct_8_components(code_2, ii);
415+
simd8float32 xi3 = quant.reconstruct_8_components(code_3, ii);
416+
sim0.add_8_components(xi0);
417+
sim1.add_8_components(xi1);
418+
sim2.add_8_components(xi2);
419+
sim3.add_8_components(xi3);
420+
}
421+
422+
dis0 = sim0.result_8();
423+
dis1 = sim1.result_8();
424+
dis2 = sim2.result_8();
425+
dis3 = sim3.result_8();
426+
}
390427
};
391428

392429
template <class Similarity>

faiss/impl/scalar_quantizer/sq-avx512.cpp

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,42 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
368368
float query_to_code(const uint8_t* code) const final {
369369
return compute_distance(q, code);
370370
}
371+
372+
void query_to_codes_batch_4(
373+
const uint8_t* code_0,
374+
const uint8_t* code_1,
375+
const uint8_t* code_2,
376+
const uint8_t* code_3,
377+
float& dis0,
378+
float& dis1,
379+
float& dis2,
380+
float& dis3) const final {
381+
Similarity sim0(q);
382+
Similarity sim1(q);
383+
Similarity sim2(q);
384+
Similarity sim3(q);
385+
386+
sim0.begin_16();
387+
sim1.begin_16();
388+
sim2.begin_16();
389+
sim3.begin_16();
390+
391+
for (size_t i = 0; i < quant.d; i += 16) {
392+
simd16float32 xi0 = quant.reconstruct_16_components(code_0, i);
393+
simd16float32 xi1 = quant.reconstruct_16_components(code_1, i);
394+
simd16float32 xi2 = quant.reconstruct_16_components(code_2, i);
395+
simd16float32 xi3 = quant.reconstruct_16_components(code_3, i);
396+
sim0.add_16_components(xi0);
397+
sim1.add_16_components(xi1);
398+
sim2.add_16_components(xi2);
399+
sim3.add_16_components(xi3);
400+
}
401+
402+
dis0 = sim0.result_16();
403+
dis1 = sim1.result_16();
404+
dis2 = sim2.result_16();
405+
dis3 = sim3.result_16();
406+
}
371407
};
372408

373409
template <class Similarity>

faiss/impl/scalar_quantizer/sq-neon.cpp

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -358,6 +358,42 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
358358
float query_to_code(const uint8_t* code) const final {
359359
return compute_distance(q, code);
360360
}
361+
362+
void query_to_codes_batch_4(
363+
const uint8_t* code_0,
364+
const uint8_t* code_1,
365+
const uint8_t* code_2,
366+
const uint8_t* code_3,
367+
float& dis0,
368+
float& dis1,
369+
float& dis2,
370+
float& dis3) const final {
371+
Similarity sim0(q);
372+
Similarity sim1(q);
373+
Similarity sim2(q);
374+
Similarity sim3(q);
375+
376+
sim0.begin_8();
377+
sim1.begin_8();
378+
sim2.begin_8();
379+
sim3.begin_8();
380+
381+
for (size_t i = 0; i < quant.d; i += 8) {
382+
simd8float32 xi0 = quant.reconstruct_8_components(code_0, i);
383+
simd8float32 xi1 = quant.reconstruct_8_components(code_1, i);
384+
simd8float32 xi2 = quant.reconstruct_8_components(code_2, i);
385+
simd8float32 xi3 = quant.reconstruct_8_components(code_3, i);
386+
sim0.add_8_components(xi0);
387+
sim1.add_8_components(xi1);
388+
sim2.add_8_components(xi2);
389+
sim3.add_8_components(xi3);
390+
}
391+
392+
dis0 = sim0.result_8();
393+
dis1 = sim1.result_8();
394+
dis2 = sim2.result_8();
395+
dis3 = sim3.result_8();
396+
}
361397
};
362398

363399
template <class Similarity>

0 commit comments

Comments
 (0)