mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
1384484b7d
* #2039 clean redundant compile options Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * #2039 add distances_avx.h and distances_avx.cpp Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * #2039 add ScalarQuantizerDC_avx.cpp, ScalarQuantizerDC_avx.h and ScalarQuantizerCodec_avx.h Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * #2039 remove -mavx2 from default compile option Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * update changelog Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * #2039 clean code Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * #2039 clean code Signed-off-by: yudong.cai <yudong.cai@zilliz.com>
105 lines
3.1 KiB
C++
105 lines
3.1 KiB
C++
|
|
// -*- c++ -*-
|
|
|
|
#include <iostream>
|
|
#include <mutex>
|
|
|
|
#include <faiss/FaissHook.h>
|
|
#include <faiss/impl/FaissAssert.h>
|
|
#include <faiss/impl/ScalarQuantizerDC.h>
|
|
#include <faiss/impl/ScalarQuantizerDC_avx.h>
|
|
#include <faiss/impl/ScalarQuantizerDC_avx512.h>
|
|
#include <faiss/utils/distances.h>
|
|
#include <faiss/utils/distances_avx.h>
|
|
#include <faiss/utils/distances_avx512.h>
|
|
#include <faiss/utils/instruction_set.h>
|
|
|
|
namespace faiss {
|
|
|
|
bool faiss_use_avx512 = true;
|
|
|
|
/* set default to AVX */
|
|
fvec_func_ptr fvec_inner_product = fvec_inner_product_avx;
|
|
fvec_func_ptr fvec_L2sqr = fvec_L2sqr_avx;
|
|
fvec_func_ptr fvec_L1 = fvec_L1_avx;
|
|
fvec_func_ptr fvec_Linf = fvec_Linf_avx;
|
|
|
|
sq_get_func_ptr sq_get_distance_computer_L2 = sq_get_distance_computer_L2_avx;
|
|
sq_get_func_ptr sq_get_distance_computer_IP = sq_get_distance_computer_IP_avx;
|
|
sq_sel_func_ptr sq_sel_quantizer = sq_select_quantizer_avx;
|
|
|
|
|
|
/*****************************************************************************/
|
|
|
|
bool support_avx512() {
|
|
if (!faiss_use_avx512) return false;
|
|
|
|
InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
|
|
return (instruction_set_inst.AVX512F() &&
|
|
instruction_set_inst.AVX512DQ() &&
|
|
instruction_set_inst.AVX512BW());
|
|
}
|
|
|
|
bool support_avx2() {
|
|
InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
|
|
return (instruction_set_inst.AVX2());
|
|
}
|
|
|
|
bool support_sse42() {
|
|
InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
|
|
return (instruction_set_inst.SSE42());
|
|
}
|
|
|
|
bool hook_init(std::string& cpu_flag) {
|
|
static std::mutex hook_mutex;
|
|
std::lock_guard<std::mutex> lock(hook_mutex);
|
|
|
|
if (support_avx512()) {
|
|
/* for IVFFLAT */
|
|
fvec_inner_product = fvec_inner_product_avx512;
|
|
fvec_L2sqr = fvec_L2sqr_avx512;
|
|
fvec_L1 = fvec_L1_avx512;
|
|
fvec_Linf = fvec_Linf_avx512;
|
|
|
|
/* for IVFSQ */
|
|
sq_get_distance_computer_L2 = sq_get_distance_computer_L2_avx512;
|
|
sq_get_distance_computer_IP = sq_get_distance_computer_IP_avx512;
|
|
sq_sel_quantizer = sq_select_quantizer_avx512;
|
|
|
|
cpu_flag = "AVX512";
|
|
} else if (support_avx2()) {
|
|
/* for IVFFLAT */
|
|
fvec_inner_product = fvec_inner_product_avx;
|
|
fvec_L2sqr = fvec_L2sqr_avx;
|
|
fvec_L1 = fvec_L1_avx;
|
|
fvec_Linf = fvec_Linf_avx;
|
|
|
|
/* for IVFSQ */
|
|
sq_get_distance_computer_L2 = sq_get_distance_computer_L2_avx;
|
|
sq_get_distance_computer_IP = sq_get_distance_computer_IP_avx;
|
|
sq_sel_quantizer = sq_select_quantizer_avx;
|
|
|
|
cpu_flag = "AVX2";
|
|
} else if (support_sse42()) {
|
|
/* for IVFFLAT */
|
|
fvec_inner_product = fvec_inner_product_sse;
|
|
fvec_L2sqr = fvec_L2sqr_sse;
|
|
fvec_L1 = fvec_L1_sse;
|
|
fvec_Linf = fvec_Linf_sse;
|
|
|
|
/* for IVFSQ */
|
|
sq_get_distance_computer_L2 = sq_get_distance_computer_L2_sse;
|
|
sq_get_distance_computer_IP = sq_get_distance_computer_IP_sse;
|
|
sq_sel_quantizer = sq_select_quantizer_sse;
|
|
|
|
cpu_flag = "SSE42";
|
|
} else {
|
|
cpu_flag = "UNSUPPORTED";
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
} // namespace faiss
|