Unverified Commit c3bf2704 authored by 蔡宇东's avatar 蔡宇东 Committed by GitHub
Browse files

support avx512 (#1195) 



* #1122 support AVX-512 for IVFFLAT distance calculation

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 support AVX-512 for IVFSQ8

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix avx512 build option

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 remove global faiss compile option -march=skylake-avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add InstructionSet and unittest

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 update InstructionSet unittest

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add FaissHook

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add distances_simd_avx512.cpp

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix build issue

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add ScalarQuantizer_avx512.cpp and IndexScalarQuantizer_avx512.cpp

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix build issue

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix clang format

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 move test_instructionset.cpp to index/unittest

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix clang format

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 set hook default to AVX

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add hidden config use_avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix faiss clone issue

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix faiss build issue

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix faiss ScalarQuantizer_avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix faiss ScalarQuantizer_avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 clean code

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 remove ScalarQuantizer_avx512 and IndexScalarQuantizer_avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 change back index_factory.cpp/index_read.cpp/index_write.cpp

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 split ScalarQuantizer to ScalarQuantizerCodec and Similarity

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 #1122 split ScalarQuantizer to ScalarQuantizerCodec

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 split ScalarQuantizerCodec.cpp to ScalarQuantizerOp.cpp

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 ScalarQuantizer support avx512

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add hook sse for ScalarQuantizer

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 correct head file

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 add changelog

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix CodeFactor

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>

* #1122 fix unittest

Signed-off-by: default avataryudong.cai <yudong.cai@zilliz.com>
parent 79ac5443

CHANGELOG.md

+2 −0
Original line number Diff line number Diff line
@@ -41,7 +41,9 @@ Please mark all change in change log and use the issue from GitHub 
- \#823 - Support binary vector tanimoto/jaccard/hamming metric

- \#853 - Support HNSW

- \#910 - Change Milvus c++ standard to c++17

- \#1122 - Support AVX-512 in FAISS

- \#1204 - Add api to get table data information

- \#1250 - Support CPU profiling

- \#1302 - Get all record IDs in a segment by given a segment id



## Improvement

core/src/index/thirdparty/faiss/FaissHook.cpp

0 → 100644
+99 −0
Original line number Diff line number Diff line 


// -*- c++ -*-



#include <iostream>

#include <mutex>



#include <faiss/FaissHook.h>

#include <faiss/impl/FaissAssert.h>

#include <faiss/impl/ScalarQuantizerDC.h>

#include <faiss/impl/ScalarQuantizerDC_avx512.h>

#include <faiss/utils/distances.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_avx() {

    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();

    return (instruction_set_inst.AVX2());

}



bool support_sse() {

    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();

    return (instruction_set_inst.SSE());

}



void hook_init() {

    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;



        std::cout << "FAISS hook AVX512" << std::endl;

    } else if (support_avx()) {

        /* 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;



        std::cout << "FAISS hook AVX" << std::endl;

    } else if (support_sse()) {

        /* 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;



        std::cout << "FAISS hook SSE" << std::endl;

    } else {

        FAISS_ASSERT_MSG(false, "CPU not supported!");

    }

}



} // namespace faiss

core/src/index/thirdparty/faiss/FaissHook.h

0 → 100644
+33 −0
Original line number Diff line number Diff line 


// -*- c++ -*-



#pragma once



#include <vector>

#include <stddef.h>

#include <faiss/impl/ScalarQuantizerOp.h>



namespace faiss {



typedef float (*fvec_func_ptr)(const float*, const float*, size_t);



typedef SQDistanceComputer* (*sq_get_func_ptr)(QuantizerType, size_t, const std::vector<float>&);

typedef Quantizer* (*sq_sel_func_ptr)(QuantizerType, size_t, const std::vector<float>&);





extern bool faiss_use_avx512;



extern fvec_func_ptr fvec_inner_product;

extern fvec_func_ptr fvec_L2sqr;

extern fvec_func_ptr fvec_L1;

extern fvec_func_ptr fvec_Linf;



extern sq_get_func_ptr sq_get_distance_computer_L2;

extern sq_get_func_ptr sq_get_distance_computer_IP;

extern sq_sel_func_ptr sq_sel_quantizer;



extern bool support_avx512();



extern void hook_init();



} // namespace faiss

core/src/index/thirdparty/faiss/Index.cpp

+1 −0
Original line number Diff line number Diff line
@@ -12,6 +12,7 @@ 
#include <faiss/impl/AuxIndexStructures.h>

#include <faiss/impl/FaissAssert.h>

#include <faiss/utils/distances.h>

#include <faiss/FaissHook.h>



#include <cstring>

core/src/index/thirdparty/faiss/Index2Layer.cpp

+1 −1
Original line number Diff line number Diff line
@@ -27,7 +27,7 @@ 
#include <faiss/impl/AuxIndexStructures.h>

#include <faiss/IndexFlat.h>

#include <faiss/utils/distances.h>



#include <faiss/FaissHook.h>



/*

#include <faiss/utils/Heap.h>