diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py index 797c02656c..0056ba01ec 100644 --- a/contrib/torch_utils.py +++ b/contrib/torch_utils.py @@ -131,7 +131,7 @@ def torch_replace_method(the_class, name, replacement, ################################################################## def handle_torch_Index(the_class): - def torch_replacement_add(self, x): + def torch_replacement_add(self, x, numeric_type = faiss.Float32): if type(x) is np.ndarray: # forward to faiss __init__.py base method return self.add_numpy(x) @@ -139,19 +139,25 @@ def torch_replacement_add(self, x): assert type(x) is torch.Tensor n, d = x.shape assert d == self.d - x_ptr = swig_ptr_from_FloatTensor(x) + if numeric_type == faiss.Float32: + x_ptr = swig_ptr_from_FloatTensor(x) + elif numeric_type == faiss.Float16: + x_ptr = swig_ptr_from_HalfTensor(x) + else: + raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ") + if x.is_cuda: assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed' # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.add_c(n, x_ptr) + self.add_c(n, x_ptr, numeric_type) else: # CPU torch - self.add_c(n, x_ptr) + self.add_c(n, x_ptr, numeric_type) - def torch_replacement_add_with_ids(self, x, ids): + def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): if type(x) is np.ndarray: # forward to faiss __init__.py base method return self.add_with_ids_numpy(x, ids) @@ -159,7 +165,12 @@ def torch_replacement_add_with_ids(self, x, ids): assert type(x) is torch.Tensor n, d = x.shape assert d == self.d - x_ptr = swig_ptr_from_FloatTensor(x) + if numeric_type == faiss.Float32: + x_ptr = swig_ptr_from_FloatTensor(x) + elif numeric_type == faiss.Float16: + x_ptr = swig_ptr_from_HalfTensor(x) + else: + raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ") assert type(ids) is torch.Tensor assert ids.shape == (n, ), 'not same number of vectors as ids' @@ -170,10 +181,10 @@ def torch_replacement_add_with_ids(self, x, ids): # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.add_with_ids_c(n, x_ptr, ids_ptr) + self.add_with_ids_c(n, x_ptr, numeric_type, ids_ptr) else: # CPU torch - self.add_with_ids_c(n, x_ptr, ids_ptr) + self.add_with_ids_c(n, x_ptr, numeric_type, ids_ptr) def torch_replacement_assign(self, x, k, labels=None): if type(x) is np.ndarray: @@ -204,7 +215,7 @@ def torch_replacement_assign(self, x, k, labels=None): return labels - def torch_replacement_train(self, x): + def torch_replacement_train(self, x, numeric_type = faiss.Float32): if type(x) is np.ndarray: # forward to faiss __init__.py base method return self.train_numpy(x) @@ -212,21 +223,31 @@ def torch_replacement_train(self, x): assert type(x) is torch.Tensor n, d = x.shape assert d == self.d - x_ptr = swig_ptr_from_FloatTensor(x) + if numeric_type == faiss.Float32: + x_ptr = swig_ptr_from_FloatTensor(x) + elif numeric_type == faiss.Float16: + x_ptr = swig_ptr_from_HalfTensor(x) + else: + raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ") if x.is_cuda: assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed' # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.train_c(n, x_ptr) + self.train_c(n, x_ptr, numeric_type) else: # CPU torch - self.train_c(n, x_ptr) + self.train_c(n, x_ptr, numeric_type) - def search_methods_common(x, k, D, I): + def search_methods_common(x, k, D, I, numeric_type=faiss.Float32): n, d = x.shape - x_ptr = swig_ptr_from_FloatTensor(x) + if numeric_type == faiss.Float32: + x_ptr = swig_ptr_from_FloatTensor(x) + elif numeric_type == faiss.Float16: + x_ptr = swig_ptr_from_HalfTensor(x) + else: + raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ") if D is None: D = torch.empty(n, k, device=x.device, dtype=torch.float32) @@ -244,7 +265,7 @@ def search_methods_common(x, k, D, I): return x_ptr, D_ptr, I_ptr, D, I - def torch_replacement_search(self, x, k, D=None, I=None): + def torch_replacement_search(self, x, k, D=None, I=None, numeric_type=faiss.Float32): if type(x) is np.ndarray: # forward to faiss __init__.py base method return self.search_numpy(x, k, D=D, I=I) @@ -260,10 +281,10 @@ def torch_replacement_search(self, x, k, D=None, I=None): # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.search_c(n, x_ptr, k, D_ptr, I_ptr) + self.search_c(n, x_ptr, numeric_type, k, D_ptr, I_ptr) else: # CPU torch - self.search_c(n, x_ptr, k, D_ptr, I_ptr) + self.search_c(n, x_ptr, numeric_type, k, D_ptr, I_ptr) return D, I diff --git a/faiss/Index.h b/faiss/Index.h index 95af05df74..0af35cfce2 100644 --- a/faiss/Index.h +++ b/faiss/Index.h @@ -61,6 +61,7 @@ struct DistanceComputer; enum NumericType { Float32, Float16, + UInt8, }; inline size_t get_numeric_type_size(NumericType numeric_type) { diff --git a/faiss/IndexAdditiveQuantizer.cpp b/faiss/IndexAdditiveQuantizer.cpp index f9e7c773e9..12a5877350 100644 --- a/faiss/IndexAdditiveQuantizer.cpp +++ b/faiss/IndexAdditiveQuantizer.cpp @@ -419,6 +419,13 @@ void AdditiveCoarseQuantizer::add(idx_t, const float*) { FAISS_THROW_MSG("not applicable"); } +void AdditiveCoarseQuantizer::add( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::add(n, x, numeric_type); +}; + void AdditiveCoarseQuantizer::reconstruct(idx_t key, float* recons) const { aq->decode_64bit(key, recons); } @@ -454,6 +461,13 @@ void AdditiveCoarseQuantizer::train(idx_t n, const float* x) { } } +void AdditiveCoarseQuantizer::train( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + void AdditiveCoarseQuantizer::search( idx_t n, const float* x, @@ -472,6 +486,17 @@ void AdditiveCoarseQuantizer::search( } } +void AdditiveCoarseQuantizer::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + /************************************************************************************** * ResidualCoarseQuantizer **************************************************************************************/ diff --git a/faiss/IndexAdditiveQuantizer.h b/faiss/IndexAdditiveQuantizer.h index 31e3c8c0a4..6f03c42ccd 100644 --- a/faiss/IndexAdditiveQuantizer.h +++ b/faiss/IndexAdditiveQuantizer.h @@ -171,6 +171,7 @@ struct AdditiveCoarseQuantizer : Index { /// N/A void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -179,9 +180,19 @@ struct AdditiveCoarseQuantizer : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void reconstruct(idx_t key, float* recons) const override; + void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// N/A void reset() override; diff --git a/faiss/IndexBinary.h b/faiss/IndexBinary.h index e9801a7db4..d1725921f7 100644 --- a/faiss/IndexBinary.h +++ b/faiss/IndexBinary.h @@ -8,6 +8,7 @@ #ifndef FAISS_INDEX_BINARY_H #define FAISS_INDEX_BINARY_H +#include #include #include #include @@ -54,6 +55,13 @@ struct IndexBinary { * @param x training vecors, size n * d / 8 */ virtual void train(idx_t n, const uint8_t* x); + virtual void train(idx_t n, const void* x, NumericType numeric_type) { + if (numeric_type == NumericType::UInt8) { + train(n, static_cast(x)); + } else { + FAISS_THROW_MSG("IndexBinary::train: unsupported numeric type"); + } + }; /** Add n vectors of dimension d to the index. * @@ -61,6 +69,13 @@ struct IndexBinary { * @param x input matrix, size n * d / 8 */ virtual void add(idx_t n, const uint8_t* x) = 0; + virtual void add(idx_t n, const void* x, NumericType numeric_type) { + if (numeric_type == NumericType::UInt8) { + add(n, static_cast(x)); + } else { + FAISS_THROW_MSG("IndexBinary::add: unsupported numeric type"); + } + }; /** Same as add, but stores xids instead of sequential ids. * @@ -70,6 +85,18 @@ struct IndexBinary { * @param xids if non-null, ids to store for the vectors (size n) */ virtual void add_with_ids(idx_t n, const uint8_t* x, const idx_t* xids); + virtual void add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) { + if (numeric_type == NumericType::UInt8) { + add_with_ids(n, static_cast(x), xids); + } else { + FAISS_THROW_MSG( + "IndexBinary::add_with_ids: unsupported numeric type"); + } + }; /** Query n vectors of dimension d to the index. * @@ -87,6 +114,25 @@ struct IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const = 0; + virtual void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const { + if (numeric_type == NumericType::UInt8) { + search(n, + static_cast(x), + k, + distances, + labels, + params); + } else { + FAISS_THROW_MSG("IndexBinary::search: unsupported numeric type"); + } + }; /** Query n vectors of dimension d to the index. * diff --git a/faiss/IndexBinaryFlat.cpp b/faiss/IndexBinaryFlat.cpp index bbb51d7c93..ec47b16a9a 100644 --- a/faiss/IndexBinaryFlat.cpp +++ b/faiss/IndexBinaryFlat.cpp @@ -25,6 +25,10 @@ void IndexBinaryFlat::add(idx_t n, const uint8_t* x) { ntotal += n; } +void IndexBinaryFlat::add(idx_t n, const void* x, NumericType numeric_type) { + IndexBinary::add(n, x, numeric_type); +} + void IndexBinaryFlat::reset() { xb.clear(); ntotal = 0; @@ -77,6 +81,17 @@ void IndexBinaryFlat::search( } } +void IndexBinaryFlat::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params) const { + IndexBinary::search(n, x, numeric_type, k, distances, labels, params); +} + size_t IndexBinaryFlat::remove_ids(const IDSelector& sel) { idx_t j = 0; for (idx_t i = 0; i < ntotal; i++) { diff --git a/faiss/IndexBinaryFlat.h b/faiss/IndexBinaryFlat.h index 0ce43f3e9d..525df46852 100644 --- a/faiss/IndexBinaryFlat.h +++ b/faiss/IndexBinaryFlat.h @@ -36,6 +36,7 @@ struct IndexBinaryFlat : IndexBinary { explicit IndexBinaryFlat(idx_t d); void add(idx_t n, const uint8_t* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -46,6 +47,14 @@ struct IndexBinaryFlat : IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void range_search( idx_t n, diff --git a/faiss/IndexBinaryHNSW.cpp b/faiss/IndexBinaryHNSW.cpp index c14ae66987..54304a3fd9 100644 --- a/faiss/IndexBinaryHNSW.cpp +++ b/faiss/IndexBinaryHNSW.cpp @@ -196,6 +196,10 @@ void IndexBinaryHNSW::train(idx_t n, const uint8_t* x) { is_trained = true; } +void IndexBinaryHNSW::train(idx_t n, const void* x, NumericType numeric_type) { + IndexBinary::train(n, x, numeric_type); +} + void IndexBinaryHNSW::search( idx_t n, const uint8_t* x, @@ -235,6 +239,17 @@ void IndexBinaryHNSW::search( } } +void IndexBinaryHNSW::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params) const { + IndexBinary::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexBinaryHNSW::add(idx_t n, const uint8_t* x) { FAISS_THROW_IF_NOT(is_trained); int n0 = ntotal; @@ -244,6 +259,10 @@ void IndexBinaryHNSW::add(idx_t n, const uint8_t* x) { hnsw_add_vertices(*this, n0, n, x, verbose, hnsw.levels.size() == ntotal); } +void IndexBinaryHNSW::add(idx_t n, const void* x, NumericType numeric_type) { + IndexBinary::add(n, x, numeric_type); +} + void IndexBinaryHNSW::reset() { hnsw.reset(); storage->reset(); diff --git a/faiss/IndexBinaryHNSW.h b/faiss/IndexBinaryHNSW.h index e2945fc10b..aa4b78bb7a 100644 --- a/faiss/IndexBinaryHNSW.h +++ b/faiss/IndexBinaryHNSW.h @@ -49,9 +49,11 @@ struct IndexBinaryHNSW : IndexBinary { DistanceComputer* get_distance_computer() const; void add(idx_t n, const uint8_t* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// Trains the storage if needed void train(idx_t n, const uint8_t* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// entry point for search void search( @@ -61,6 +63,14 @@ struct IndexBinaryHNSW : IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void reconstruct(idx_t key, uint8_t* recons) const override; diff --git a/faiss/IndexFastScan.cpp b/faiss/IndexFastScan.cpp index b18d15bc17..3d55f7fe56 100644 --- a/faiss/IndexFastScan.cpp +++ b/faiss/IndexFastScan.cpp @@ -94,6 +94,10 @@ void IndexFastScan::add(idx_t n, const float* x) { ntotal += n; } +void IndexFastScan::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + CodePacker* IndexFastScan::get_CodePacker() const { return new CodePackerPQ4(M, bbs); } @@ -270,6 +274,17 @@ void IndexFastScan::search( } } +void IndexFastScan::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + template void IndexFastScan::search_dispatch_implem( idx_t n, diff --git a/faiss/IndexFastScan.h b/faiss/IndexFastScan.h index a0f5c592f0..ac06c6eb6f 100644 --- a/faiss/IndexFastScan.h +++ b/faiss/IndexFastScan.h @@ -72,8 +72,17 @@ struct IndexFastScan : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; virtual void compute_codes(uint8_t* codes, idx_t n, const float* x) const = 0; diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index 84c33f970c..aad7a662af 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -56,6 +56,17 @@ void IndexFlat::search( } } +void IndexFlat::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexFlat::range_search( idx_t n, const float* x, @@ -404,6 +415,10 @@ void IndexFlat1D::add(idx_t n, const float* x) { update_permutation(); } +void IndexFlat1D::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void IndexFlat1D::reset() { IndexFlatL2::reset(); perm.clear(); @@ -518,4 +533,15 @@ void IndexFlat1D::search( } } +void IndexFlat1D::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + } // namespace faiss diff --git a/faiss/IndexFlat.h b/faiss/IndexFlat.h index 876a2ec2fa..5c994860bc 100644 --- a/faiss/IndexFlat.h +++ b/faiss/IndexFlat.h @@ -29,6 +29,14 @@ struct IndexFlat : IndexFlatCodes { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void range_search( idx_t n, @@ -112,6 +120,7 @@ struct IndexFlat1D : IndexFlatL2 { void update_permutation(); void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -123,6 +132,14 @@ struct IndexFlat1D : IndexFlatL2 { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; }; } // namespace faiss diff --git a/faiss/IndexFlatCodes.cpp b/faiss/IndexFlatCodes.cpp index d5b86b385e..ab7fd1abe1 100644 --- a/faiss/IndexFlatCodes.cpp +++ b/faiss/IndexFlatCodes.cpp @@ -32,6 +32,10 @@ void IndexFlatCodes::add(idx_t n, const float* x) { ntotal += n; } +void IndexFlatCodes::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +}; + void IndexFlatCodes::add_sa_codes( idx_t n, const uint8_t* codes_in, @@ -266,6 +270,17 @@ void IndexFlatCodes::search( n, distances, labels, k, metric_type, sel, r, this, x); } +void IndexFlatCodes::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexFlatCodes::range_search( idx_t n, const float* x, diff --git a/faiss/IndexFlatCodes.h b/faiss/IndexFlatCodes.h index 56a11df795..dd0ee8f4ca 100644 --- a/faiss/IndexFlatCodes.h +++ b/faiss/IndexFlatCodes.h @@ -31,6 +31,7 @@ struct IndexFlatCodes : Index { /// default add uses sa_encode void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -63,6 +64,14 @@ struct IndexFlatCodes : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void range_search( idx_t n, diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index 1ee15f4484..58ad28300f 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -230,6 +230,10 @@ void IndexHNSW::train(idx_t n, const float* x) { is_trained = true; } +void IndexHNSW::train(idx_t n, const void* x, NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + namespace { template @@ -311,6 +315,17 @@ void IndexHNSW::search( } } +void IndexHNSW::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexHNSW::range_search( idx_t n, const float* x, @@ -342,6 +357,10 @@ void IndexHNSW::add(idx_t n, const float* x) { hnsw_add_vertices(*this, n0, n, x, verbose, hnsw.levels.size() == ntotal); } +void IndexHNSW::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void IndexHNSW::reset() { hnsw.reset(); storage->reset(); @@ -657,6 +676,10 @@ void IndexHNSWPQ::train(idx_t n, const float* x) { (dynamic_cast(storage))->pq.compute_sdc_table(); } +void IndexHNSWPQ::train(idx_t n, const void* x, NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + /************************************************************** * IndexHNSWSQ implementation **************************************************************/ @@ -864,6 +887,17 @@ void IndexHNSW2Level::search( } } +void IndexHNSW2Level::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexHNSW2Level::flip_to_ivf() { Index2Layer* storage2l = dynamic_cast(storage); @@ -933,6 +967,10 @@ void IndexHNSWCagra::add(idx_t n, const float* x) { IndexHNSW::add(n, x); } +void IndexHNSWCagra::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void IndexHNSWCagra::search( idx_t n, const float* x, @@ -984,6 +1022,17 @@ void IndexHNSWCagra::search( } } +void IndexHNSWCagra::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + faiss::NumericType IndexHNSWCagra::get_numeric_type() const { return numeric_type_; } diff --git a/faiss/IndexHNSW.h b/faiss/IndexHNSW.h index c6e80df462..df621af934 100644 --- a/faiss/IndexHNSW.h +++ b/faiss/IndexHNSW.h @@ -53,9 +53,11 @@ struct IndexHNSW : Index { ~IndexHNSW() override; void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// Trains the storage if needed void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// entry point for search void search( @@ -65,6 +67,14 @@ struct IndexHNSW : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void range_search( idx_t n, @@ -137,6 +147,7 @@ struct IndexHNSWPQ : IndexHNSW { int pq_nbits = 8, MetricType metric = METRIC_L2); void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; }; /** SQ index topped with a HNSW structure to access elements @@ -167,6 +178,14 @@ struct IndexHNSW2Level : IndexHNSW { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; }; struct IndexHNSWCagra : IndexHNSW { @@ -191,6 +210,7 @@ struct IndexHNSWCagra : IndexHNSW { int num_base_level_search_entrypoints = 32; void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// entry point for search void search( @@ -200,6 +220,14 @@ struct IndexHNSWCagra : IndexHNSW { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; faiss::NumericType get_numeric_type() const; void set_numeric_type(faiss::NumericType numeric_type); diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index 6f4530544a..753254da04 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -12,6 +12,7 @@ #include #include #include +#include "faiss/Index.h" #include #include @@ -33,6 +34,17 @@ void sync_d(IndexBinary* index) { } // anonymous namespace +template +NumericType component_t_to_numeric() { + if constexpr (std::is_same::value) { + return NumericType::Float32; + } else if constexpr (std::is_same::value) { + return NumericType::UInt8; + } else { + FAISS_THROW_MSG("Unsupported component_t"); + } +} + /***************************************************** * IndexIDMap implementation *******************************************************/ @@ -47,6 +59,16 @@ IndexIDMapTemplate::IndexIDMapTemplate(IndexT* index) : index(index) { sync_d(this); } +template +void IndexIDMapTemplate::add( + idx_t, + const void*, + NumericType numeric_type) { + FAISS_THROW_MSG( + "add does not make sense with IndexIDMap, " + "use add_with_ids"); +} + template void IndexIDMapTemplate::add( idx_t, @@ -59,11 +81,21 @@ void IndexIDMapTemplate::add( template void IndexIDMapTemplate::train( idx_t n, - const typename IndexT::component_t* x) { - index->train(n, x); + const void* x, + NumericType numeric_type) { + index->train(n, x, numeric_type); this->is_trained = index->is_trained; } +template +void IndexIDMapTemplate::train( + idx_t n, + const typename IndexT::component_t* x) { + train(n, + static_cast(x), + component_t_to_numeric()); +} + template void IndexIDMapTemplate::reset() { index->reset(); @@ -74,14 +106,27 @@ void IndexIDMapTemplate::reset() { template void IndexIDMapTemplate::add_with_ids( idx_t n, - const typename IndexT::component_t* x, + const void* x, + NumericType numeric_type, const idx_t* xids) { - index->add(n, x); + index->add(n, x, numeric_type); for (idx_t i = 0; i < n; i++) id_map.push_back(xids[i]); this->ntotal = index->ntotal; } +template +void IndexIDMapTemplate::add_with_ids( + idx_t n, + const typename IndexT::component_t* x, + const idx_t* xids) { + add_with_ids( + n, + static_cast(x), + component_t_to_numeric(), + xids); +} + template size_t IndexIDMapTemplate::sa_code_size() const { return index->sa_code_size(); @@ -123,7 +168,8 @@ struct ScopedSelChange { template void IndexIDMapTemplate::search( idx_t n, - const typename IndexT::component_t* x, + const void* x, + NumericType numeric_type, idx_t k, typename IndexT::distance_t* distances, idx_t* labels, @@ -147,7 +193,7 @@ void IndexIDMapTemplate::search( sel_change.set(params_non_const, &this_idtrans); } } - index->search(n, x, k, distances, labels, params); + index->search(n, x, numeric_type, k, distances, labels, params); idx_t* li = labels; #pragma omp parallel for for (idx_t i = 0; i < n * k; i++) { @@ -155,6 +201,23 @@ void IndexIDMapTemplate::search( } } +template +void IndexIDMapTemplate::search( + idx_t n, + const typename IndexT::component_t* x, + idx_t k, + typename IndexT::distance_t* distances, + idx_t* labels, + const SearchParameters* params) const { + search(n, + static_cast(x), + component_t_to_numeric(), + k, + distances, + labels, + params); +} + template void IndexIDMapTemplate::range_search( idx_t n, @@ -238,15 +301,28 @@ IndexIDMap2Template::IndexIDMap2Template(IndexT* index) template void IndexIDMap2Template::add_with_ids( idx_t n, - const typename IndexT::component_t* x, + const void* x, + NumericType numeric_type, const idx_t* xids) { size_t prev_ntotal = this->ntotal; - IndexIDMapTemplate::add_with_ids(n, x, xids); + IndexIDMapTemplate::add_with_ids(n, x, numeric_type, xids); for (size_t i = prev_ntotal; i < this->ntotal; i++) { rev_map[this->id_map[i]] = i; } } +template +void IndexIDMap2Template::add_with_ids( + idx_t n, + const typename IndexT::component_t* x, + const idx_t* xids) { + add_with_ids( + n, + static_cast(x), + component_t_to_numeric(), + xids); +} + template void IndexIDMap2Template::check_consistency() const { FAISS_THROW_IF_NOT(rev_map.size() == this->id_map.size()); diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h index dd3887ae76..cfe43377ee 100644 --- a/faiss/IndexIDMap.h +++ b/faiss/IndexIDMap.h @@ -31,9 +31,15 @@ struct IndexIDMapTemplate : IndexT { /// @param xids if non-null, ids to store for the vectors (size n) void add_with_ids(idx_t n, const component_t* x, const idx_t* xids) override; + void add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) override; /// this will fail. Use add_with_ids void add(idx_t n, const component_t* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -42,8 +48,17 @@ struct IndexIDMapTemplate : IndexT { distance_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + distance_t* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void train(idx_t n, const component_t* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -89,6 +104,11 @@ struct IndexIDMap2Template : IndexIDMapTemplate { void add_with_ids(idx_t n, const component_t* x, const idx_t* xids) override; + void add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) override; size_t remove_ids(const IDSelector& sel) override; diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index 7c775363b8..fdeae477db 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -181,12 +181,24 @@ void IndexIVF::add(idx_t n, const float* x) { add_with_ids(n, x, nullptr); } +void IndexIVF::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void IndexIVF::add_with_ids(idx_t n, const float* x, const idx_t* xids) { std::unique_ptr coarse_idx(new idx_t[n]); quantizer->assign(n, x, coarse_idx.get()); add_core(n, x, xids, coarse_idx.get()); } +void IndexIVF::add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) { + Index::add_with_ids(n, x, numeric_type, xids); +} + void IndexIVF::add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids) { size_t coarse_size = coarse_code_size(); DirectMapAdd dm_adder(direct_map, n, xids); @@ -389,6 +401,17 @@ void IndexIVF::search( } } +void IndexIVF::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexIVF::search_preassigned( idx_t n, const float* x, @@ -1182,6 +1205,10 @@ void IndexIVF::train(idx_t n, const float* x) { is_trained = true; } +void IndexIVF::train(idx_t n, const void* x, NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + idx_t IndexIVF::train_encoder_num_vectors() const { return 0; } diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index 304b9d1bdb..c21fb53c99 100644 --- a/faiss/IndexIVF.h +++ b/faiss/IndexIVF.h @@ -217,12 +217,19 @@ struct IndexIVF : Index, IndexIVFInterface { /// Trains the quantizer and calls train_encoder to train sub-quantizers void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// Calls add_with_ids with NULL ids void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// default implementation that calls encode_vectors void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; + void add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) override; /** Implementation of vector addition where the vector assignments are * predefined. The default implementation hands over the code extraction to @@ -317,6 +324,14 @@ struct IndexIVF : Index, IndexIVFInterface { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void range_search( idx_t n, diff --git a/faiss/IndexIVFFlat.cpp b/faiss/IndexIVFFlat.cpp index 95fd7bc0bb..9ac21b4f17 100644 --- a/faiss/IndexIVFFlat.cpp +++ b/faiss/IndexIVFFlat.cpp @@ -292,6 +292,13 @@ void IndexIVFFlatDedup::train(idx_t n, const float* x) { IndexIVFFlat::train(n2, x2.get()); } +void IndexIVFFlatDedup::train( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + void IndexIVFFlatDedup::add_with_ids( idx_t na, const float* x, @@ -360,6 +367,14 @@ void IndexIVFFlatDedup::add_with_ids( ntotal += n_add; } +void IndexIVFFlatDedup::add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) { + Index::add_with_ids(n, x, numeric_type, xids); +} + void IndexIVFFlatDedup::search_preassigned( idx_t n, const float* x, diff --git a/faiss/IndexIVFFlat.h b/faiss/IndexIVFFlat.h index a25f8bacf5..0df24eae5f 100644 --- a/faiss/IndexIVFFlat.h +++ b/faiss/IndexIVFFlat.h @@ -77,9 +77,15 @@ struct IndexIVFFlatDedup : IndexIVFFlat { /// also dedups the training set void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// implemented for all IndexIVF* classes void add_with_ids(idx_t n, const float* x, const idx_t* xids) override; + void add_with_ids( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) override; void search_preassigned( idx_t n, diff --git a/faiss/IndexNNDescent.cpp b/faiss/IndexNNDescent.cpp index 696a979b39..6ab9b7de05 100644 --- a/faiss/IndexNNDescent.cpp +++ b/faiss/IndexNNDescent.cpp @@ -100,6 +100,10 @@ void IndexNNDescent::train(idx_t n, const float* x) { is_trained = true; } +void IndexNNDescent::train(idx_t n, const void* x, NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + void IndexNNDescent::search( idx_t n, const float* x, @@ -152,6 +156,17 @@ void IndexNNDescent::search( } } +void IndexNNDescent::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexNNDescent::add(idx_t n, const float* x) { FAISS_THROW_IF_NOT_MSG( storage, @@ -172,6 +187,10 @@ void IndexNNDescent::add(idx_t n, const float* x) { nndescent.build(*dis, ntotal, verbose); } +void IndexNNDescent::add(idx_t n, const void* x, NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void IndexNNDescent::reset() { nndescent.reset(); storage->reset(); diff --git a/faiss/IndexNNDescent.h b/faiss/IndexNNDescent.h index 0de302b586..ae4d69e4ac 100644 --- a/faiss/IndexNNDescent.h +++ b/faiss/IndexNNDescent.h @@ -42,9 +42,11 @@ struct IndexNNDescent : Index { ~IndexNNDescent() override; void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// Trains the storage if needed void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// entry point for search void search( @@ -54,7 +56,14 @@ struct IndexNNDescent : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; - + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void reconstruct(idx_t key, float* recons) const override; void reset() override; diff --git a/faiss/IndexPQ.cpp b/faiss/IndexPQ.cpp index 8193e78b17..b645afd544 100644 --- a/faiss/IndexPQ.cpp +++ b/faiss/IndexPQ.cpp @@ -71,6 +71,10 @@ void IndexPQ::train(idx_t n, const float* x) { is_trained = true; } +void IndexPQ::train(idx_t n, const void* x, NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + namespace { template @@ -256,6 +260,17 @@ void IndexPQ::search( } } +void IndexPQ::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void IndexPQStats::reset() { nq = ncode = n_hamming_pass = 0; } @@ -874,6 +889,13 @@ void MultiIndexQuantizer::train(idx_t n, const float* x) { ntotal *= pq.ksub; } +void MultiIndexQuantizer::train( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + // block size used in MultiIndexQuantizer::search int multi_index_quantizer_search_bs = 32768; @@ -956,6 +978,17 @@ void MultiIndexQuantizer::search( } } +void MultiIndexQuantizer::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + void MultiIndexQuantizer::reconstruct(idx_t key, float* recons) const { int64_t jj = key; for (int m = 0; m < pq.M; m++) { @@ -972,6 +1005,13 @@ void MultiIndexQuantizer::add(idx_t /*n*/, const float* /*x*/) { "it does not support add"); } +void MultiIndexQuantizer::add( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::add(n, x, numeric_type); +} + void MultiIndexQuantizer::reset() { FAISS_THROW_MSG( "This index has virtual elements, " @@ -1021,6 +1061,13 @@ void MultiIndexQuantizer2::train(idx_t n, const float* x) { } } +void MultiIndexQuantizer2::train( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + void MultiIndexQuantizer2::search( idx_t n, const float* x, @@ -1112,4 +1159,15 @@ void MultiIndexQuantizer2::search( } } +void MultiIndexQuantizer2::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + } // namespace faiss diff --git a/faiss/IndexPQ.h b/faiss/IndexPQ.h index 2954f580f0..e8cfc27c7a 100644 --- a/faiss/IndexPQ.h +++ b/faiss/IndexPQ.h @@ -36,6 +36,7 @@ struct IndexPQ : IndexFlatCodes { IndexPQ(); void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -44,6 +45,14 @@ struct IndexPQ : IndexFlatCodes { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; /* The standalone codec interface */ void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override; @@ -142,6 +151,7 @@ struct MultiIndexQuantizer : Index { size_t nbits); ///< number of bit per subvector index void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -150,9 +160,19 @@ struct MultiIndexQuantizer : Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; /// add and reset will crash at runtime void add(idx_t n, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; + void reset() override; MultiIndexQuantizer() {} @@ -179,6 +199,7 @@ struct MultiIndexQuantizer2 : MultiIndexQuantizer { Index* assign_index_1); void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -187,6 +208,14 @@ struct MultiIndexQuantizer2 : MultiIndexQuantizer { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; }; } // namespace faiss diff --git a/faiss/IndexScalarQuantizer.cpp b/faiss/IndexScalarQuantizer.cpp index cc8d939e34..6e606bfce8 100644 --- a/faiss/IndexScalarQuantizer.cpp +++ b/faiss/IndexScalarQuantizer.cpp @@ -45,6 +45,13 @@ void IndexScalarQuantizer::train(idx_t n, const float* x) { is_trained = true; } +void IndexScalarQuantizer::train( + idx_t n, + const void* x, + NumericType numeric_type) { + Index::train(n, x, numeric_type); +} + void IndexScalarQuantizer::search( idx_t n, const float* x, @@ -89,6 +96,17 @@ void IndexScalarQuantizer::search( } } +void IndexScalarQuantizer::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + Index::search(n, x, numeric_type, k, distances, labels, params); +} + FlatCodesDistanceComputer* IndexScalarQuantizer::get_FlatCodesDistanceComputer() const { ScalarQuantizer::SQDistanceComputer* dc = diff --git a/faiss/IndexScalarQuantizer.h b/faiss/IndexScalarQuantizer.h index e57a451065..edf24b946a 100644 --- a/faiss/IndexScalarQuantizer.h +++ b/faiss/IndexScalarQuantizer.h @@ -40,6 +40,7 @@ struct IndexScalarQuantizer : IndexFlatCodes { IndexScalarQuantizer(); void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -48,6 +49,14 @@ struct IndexScalarQuantizer : IndexFlatCodes { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override; diff --git a/faiss/gpu/GpuIndexBinaryCagra.cu b/faiss/gpu/GpuIndexBinaryCagra.cu index d3a45c29be..c321257cd4 100644 --- a/faiss/gpu/GpuIndexBinaryCagra.cu +++ b/faiss/gpu/GpuIndexBinaryCagra.cu @@ -92,10 +92,24 @@ void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) { this->ntotal = n; } +void GpuIndexBinaryCagra::train( + idx_t n, + const void* x, + NumericType numeric_type) { + IndexBinary::train(n, x, numeric_type); +} + void GpuIndexBinaryCagra::add(idx_t n, const uint8_t* x) { train(n, x); } +void GpuIndexBinaryCagra::add( + idx_t n, + const void* x, + NumericType numeric_type) { + IndexBinary::add(n, x, numeric_type); +} + void GpuIndexBinaryCagra::search( idx_t n, const uint8_t* x, @@ -158,6 +172,17 @@ void GpuIndexBinaryCagra::search( fromDevice(outIndices, labels, stream); } +void GpuIndexBinaryCagra::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int* distances, + idx_t* labels, + const SearchParameters* params) const { + IndexBinary::search(n, x, numeric_type, k, distances, labels, params); +} + void GpuIndexBinaryCagra::searchNonPaged_( idx_t n, const uint8_t* x, diff --git a/faiss/gpu/GpuIndexBinaryCagra.h b/faiss/gpu/GpuIndexBinaryCagra.h index 7951671083..5ea3ea303d 100644 --- a/faiss/gpu/GpuIndexBinaryCagra.h +++ b/faiss/gpu/GpuIndexBinaryCagra.h @@ -52,6 +52,7 @@ struct GpuIndexBinaryCagra : public IndexBinary { /// the base dataset. Use this function when you want to add vectors with /// ids. Ref: https://github.com/facebookresearch/faiss/issues/4107 void add(idx_t n, const uint8_t* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// Trains CAGRA based on the given vector data. /// NB: The use of the train function here is to build the CAGRA graph on @@ -59,6 +60,7 @@ struct GpuIndexBinaryCagra : public IndexBinary { /// of vectors (without IDs) to the index. There is no external quantizer to /// be trained here. void train(idx_t n, const uint8_t* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// Initialize ourselves from the given CPU index; will overwrite /// all data in ourselves @@ -80,6 +82,14 @@ struct GpuIndexBinaryCagra : public IndexBinary { int* distances, faiss::idx_t* labels, const faiss::SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; protected: /// Called from search when the input data is on the CPU; diff --git a/faiss/gpu/GpuIndexBinaryFlat.cu b/faiss/gpu/GpuIndexBinaryFlat.cu index b09b7220d3..8f8f07fb7d 100644 --- a/faiss/gpu/GpuIndexBinaryFlat.cu +++ b/faiss/gpu/GpuIndexBinaryFlat.cu @@ -124,6 +124,10 @@ void GpuIndexBinaryFlat::add(idx_t n, const uint8_t* x) { this->ntotal += n; } +void GpuIndexBinaryFlat::add(idx_t n, const void* x, NumericType numeric_type) { + IndexBinary::add(n, x, numeric_type); +} + void GpuIndexBinaryFlat::reset() { DeviceScope scope(binaryFlatConfig_.device); @@ -193,6 +197,17 @@ void GpuIndexBinaryFlat::search( fromDevice(outIndices, labels, stream); } +void GpuIndexBinaryFlat::search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params) const { + IndexBinary::search(n, x, numeric_type, k, distances, labels, params); +} + void GpuIndexBinaryFlat::searchNonPaged_( idx_t n, const uint8_t* x, diff --git a/faiss/gpu/GpuIndexBinaryFlat.h b/faiss/gpu/GpuIndexBinaryFlat.h index 32b76d6b47..c5d6f5f639 100644 --- a/faiss/gpu/GpuIndexBinaryFlat.h +++ b/faiss/gpu/GpuIndexBinaryFlat.h @@ -54,6 +54,7 @@ class GpuIndexBinaryFlat : public IndexBinary { void copyTo(faiss::IndexBinaryFlat* index) const; void add(faiss::idx_t n, const uint8_t* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -65,6 +66,14 @@ class GpuIndexBinaryFlat : public IndexBinary { int32_t* distances, faiss::idx_t* labels, const faiss::SearchParameters* params = nullptr) const override; + void search( + idx_t n, + const void* x, + NumericType numeric_type, + idx_t k, + int32_t* distances, + idx_t* labels, + const SearchParameters* params = nullptr) const override; void reconstruct(faiss::idx_t key, uint8_t* recons) const override; diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu index 6bc4bc1cf5..eada5936ff 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -152,6 +152,14 @@ void GpuIndexCagra::addImpl_(idx_t n, const float* x, const idx_t* ids) { FAISS_THROW_MSG("adding vectors is not supported by GpuIndexCagra."); }; +void GpuIndexCagra::addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) { + GpuIndex::addImpl_(n, x, numeric_type, ids); +} + void GpuIndexCagra::searchImpl_( idx_t n, const void* x, diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index cf4a706e7d..49e2e0e800 100644 --- a/faiss/gpu/GpuIndexCagra.h +++ b/faiss/gpu/GpuIndexCagra.h @@ -285,6 +285,11 @@ struct GpuIndexCagra : public GpuIndex { bool addImplRequiresIDs_() const override; void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + void addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) override; /// Called from GpuIndex for search void searchImpl_( diff --git a/faiss/gpu/GpuIndexFlat.cu b/faiss/gpu/GpuIndexFlat.cu index eb87e082e9..c1b03456d0 100644 --- a/faiss/gpu/GpuIndexFlat.cu +++ b/faiss/gpu/GpuIndexFlat.cu @@ -167,6 +167,10 @@ void GpuIndexFlat::train(idx_t n, const float* x) { // nothing to do } +void GpuIndexFlat::train(idx_t n, const void* x, NumericType numeric_type) { + GpuIndex::train(n, x, numeric_type); +} + void GpuIndexFlat::add(idx_t n, const float* x) { DeviceScope scope(config_.device); @@ -191,6 +195,10 @@ void GpuIndexFlat::add(idx_t n, const float* x) { } } +void GpuIndexFlat::add(idx_t n, const void* x, NumericType numeric_type) { + GpuIndex::add(n, x, numeric_type); +} + bool GpuIndexFlat::addImplRequiresIDs_() const { return false; } @@ -208,6 +216,14 @@ void GpuIndexFlat::addImpl_(idx_t n, const float* x, const idx_t* ids) { this->ntotal += n; } +void GpuIndexFlat::addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) { + GpuIndex::addImpl_(n, x, numeric_type, ids); +} + void GpuIndexFlat::searchImpl_( idx_t n, const float* x, @@ -228,6 +244,17 @@ void GpuIndexFlat::searchImpl_( queries, k, metric_type, metric_arg, outDistances, outLabels, true); } +void GpuIndexFlat::searchImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + GpuIndex::searchImpl_(n, x, numeric_type, k, distances, labels, params); +} + void GpuIndexFlat::reconstruct(idx_t key, float* out) const { DeviceScope scope(config_.device); diff --git a/faiss/gpu/GpuIndexFlat.h b/faiss/gpu/GpuIndexFlat.h index ee4c14466e..4c541f9e36 100644 --- a/faiss/gpu/GpuIndexFlat.h +++ b/faiss/gpu/GpuIndexFlat.h @@ -81,9 +81,11 @@ class GpuIndexFlat : public GpuIndex { /// This index is not trained, so this does nothing void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; /// Overrides to avoid excessive copies void add(idx_t, const float* x) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; /// Reconstruction methods; prefer the batch reconstruct as it will /// be more efficient @@ -121,6 +123,11 @@ class GpuIndexFlat : public GpuIndex { /// Called from GpuIndex for add void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + void addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) override; /// Called from GpuIndex for search void searchImpl_( @@ -130,6 +137,14 @@ class GpuIndexFlat : public GpuIndex { float* distances, idx_t* labels, const SearchParameters* params) const override; + void searchImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const override; protected: /// Our configuration options diff --git a/faiss/gpu/GpuIndexIVF.cu b/faiss/gpu/GpuIndexIVF.cu index 357c4ee77e..99ecd23748 100644 --- a/faiss/gpu/GpuIndexIVF.cu +++ b/faiss/gpu/GpuIndexIVF.cu @@ -16,6 +16,7 @@ #include #include #include +#include "GpuIndexIVF.h" namespace faiss { namespace gpu { @@ -297,6 +298,14 @@ void GpuIndexIVF::addImpl_(idx_t n, const float* x, const idx_t* xids) { ntotal += n; } +void GpuIndexIVF::addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) { + GpuIndex::addImpl_(n, x, numeric_type, ids); +} + int GpuIndexIVF::getCurrentNProbe_(const SearchParameters* params) const { size_t use_nprobe = nprobe; if (params) { @@ -345,6 +354,17 @@ void GpuIndexIVF::searchImpl_( quantizer, queries, use_nprobe, k, outDistances, outLabels); } +void GpuIndexIVF::searchImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + GpuIndex::searchImpl_(n, x, numeric_type, k, distances, labels, params); +} + void GpuIndexIVF::search_preassigned( idx_t n, const float* x, diff --git a/faiss/gpu/GpuIndexIVF.h b/faiss/gpu/GpuIndexIVF.h index d6fd5b6ffa..b6c3be5d7f 100644 --- a/faiss/gpu/GpuIndexIVF.h +++ b/faiss/gpu/GpuIndexIVF.h @@ -132,6 +132,11 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface { /// Called from GpuIndex for add/add_with_ids void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + void addImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* ids) override; /// Called from GpuIndex for search void searchImpl_( @@ -142,6 +147,15 @@ class GpuIndexIVF : public GpuIndex, public IndexIVFInterface { idx_t* labels, const SearchParameters* params) const override; + void searchImpl_( + idx_t n, + const void* x, + NumericType numeric_type, + int k, + float* distances, + idx_t* labels, + const SearchParameters* params) const override; + protected: /// Our configuration options const GpuIndexIVFConfig ivfConfig_; diff --git a/faiss/gpu/GpuIndexIVFFlat.cu b/faiss/gpu/GpuIndexIVFFlat.cu index 1266b992cf..9d2d07ecda 100644 --- a/faiss/gpu/GpuIndexIVFFlat.cu +++ b/faiss/gpu/GpuIndexIVFFlat.cu @@ -324,6 +324,10 @@ void GpuIndexIVFFlat::train(idx_t n, const float* x) { this->is_trained = true; } +void GpuIndexIVFFlat::train(idx_t n, const void* x, NumericType numeric_type) { + GpuIndex::train(n, x, numeric_type); +} + void GpuIndexIVFFlat::setIndex_( GpuResources* resources, int dim, diff --git a/faiss/gpu/GpuIndexIVFFlat.h b/faiss/gpu/GpuIndexIVFFlat.h index d6826cd9f9..2263b60064 100644 --- a/faiss/gpu/GpuIndexIVFFlat.h +++ b/faiss/gpu/GpuIndexIVFFlat.h @@ -86,6 +86,7 @@ class GpuIndexIVFFlat : public GpuIndexIVF { /// Trains the coarse quantizer based on the given vector data void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void reconstruct_n(idx_t i0, idx_t n, float* out) const override; diff --git a/faiss/gpu/GpuIndexIVFPQ.cu b/faiss/gpu/GpuIndexIVFPQ.cu index 84c4bb7957..efa58c642f 100644 --- a/faiss/gpu/GpuIndexIVFPQ.cu +++ b/faiss/gpu/GpuIndexIVFPQ.cu @@ -489,6 +489,10 @@ void GpuIndexIVFPQ::train(idx_t n, const float* x) { this->is_trained = true; } +void GpuIndexIVFPQ::train(idx_t n, const void* x, NumericType numeric_type) { + GpuIndex::train(n, x, numeric_type); +} + void GpuIndexIVFPQ::setIndex_( GpuResources* resources, int dim, diff --git a/faiss/gpu/GpuIndexIVFPQ.h b/faiss/gpu/GpuIndexIVFPQ.h index 072a0d81d5..4e73b9dc10 100644 --- a/faiss/gpu/GpuIndexIVFPQ.h +++ b/faiss/gpu/GpuIndexIVFPQ.h @@ -127,6 +127,7 @@ class GpuIndexIVFPQ : public GpuIndexIVF { /// Trains the coarse and product quantizer based on the given vector data void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; public: /// Like the CPU version, we expose a publically-visible ProductQuantizer diff --git a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu index 3c856ba5ee..68670f1974 100644 --- a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu +++ b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu @@ -279,5 +279,12 @@ void GpuIndexIVFScalarQuantizer::train(idx_t n, const float* x) { this->is_trained = true; } +void GpuIndexIVFScalarQuantizer::train( + idx_t n, + const void* x, + NumericType numeric_type) { + GpuIndex::train(n, x, numeric_type); +} + } // namespace gpu } // namespace faiss diff --git a/faiss/gpu/GpuIndexIVFScalarQuantizer.h b/faiss/gpu/GpuIndexIVFScalarQuantizer.h index 44a8c1b5a8..2ab476144a 100644 --- a/faiss/gpu/GpuIndexIVFScalarQuantizer.h +++ b/faiss/gpu/GpuIndexIVFScalarQuantizer.h @@ -88,6 +88,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF { /// Trains the coarse and scalar quantizer based on the given vector data void train(idx_t n, const float* x) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; protected: /// Validates index SQ parameters diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 9c9297c888..9fd3e3b128 100644 --- a/faiss/gpu/test/test_cagra.py +++ b/faiss/gpu/test/test_cagra.py @@ -15,7 +15,7 @@ "only if cuVS is compiled in") class TestComputeGT(unittest.TestCase): - def do_compute_GT(self, metric): + def do_compute_GT(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) @@ -31,57 +31,31 @@ def do_compute_GT(self, metric): cagraIndexConfig.build_algo = faiss.graph_build_algo_IVF_PQ index = faiss.GpuIndexCagra(res, d, metric, cagraIndexConfig) - index.train(ds.get_database()) - Dnew, Inew = index.search(ds.get_queries(), k) + database = ds.get_database().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_database() + index.train(database, numeric_type=numeric_type) + queries = ds.get_queries().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_queries() + Dnew, Inew = index.search(queries, k, numeric_type=numeric_type) evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) def test_compute_GT_L2(self): - self.do_compute_GT(faiss.METRIC_L2) + self.do_compute_GT(faiss.METRIC_L2, faiss.Float32) def test_compute_GT_IP(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT) + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float32) -@unittest.skipIf( - "CUVS" not in faiss.get_compile_options(), - "only if cuVS is compiled in") -class TestComputeGTFP16(unittest.TestCase): - - def do_compute_GT(self, metric): - d = 64 - k = 12 - ds = datasets.SyntheticDataset(d, 0, 10000, 100) - Dref, Iref = faiss.knn(ds.get_queries(), ds.get_database(), k, metric) - - res = faiss.StandardGpuResources() - - # attempt to set custom IVF-PQ params - cagraIndexConfig = faiss.GpuIndexCagraConfig() - cagraIndexIVFPQConfig = faiss.IVFPQBuildCagraConfig() - cagraIndexIVFPQConfig.kmeans_trainset_fraction = 0.1 - cagraIndexConfig.ivf_pq_params = cagraIndexIVFPQConfig - cagraIndexConfig.build_algo = faiss.graph_build_algo_IVF_PQ - - index = faiss.GpuIndexCagra(res, d, metric, cagraIndexConfig) - fp16_data = ds.get_database().astype(np.float16) - index.train(fp16_data, faiss.Float16) - fp16_queries = ds.get_queries().astype(np.float16) - Dnew, Inew = index.search(fp16_queries, k, numeric_type=faiss.Float16) - - evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) + def test_compute_GT_L2_FP16(self): + self.do_compute_GT(faiss.METRIC_L2, faiss.Float16) - def test_compute_GT_L2(self): - self.do_compute_GT(faiss.METRIC_L2) - - def test_compute_GT_IP(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT) + def test_compute_GT_IP_FP16(self): + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float16) @unittest.skipIf( "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in") class TestInterop(unittest.TestCase): - def do_interop(self, metric): + def do_interop(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) @@ -89,10 +63,13 @@ def do_interop(self, metric): res = faiss.StandardGpuResources() index = faiss.GpuIndexCagra(res, d, metric) - index.train(ds.get_database()) - Dnew, Inew = index.search(ds.get_queries(), k) + database = ds.get_database().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_database() + index.train(database, numeric_type=numeric_type) + queries = ds.get_queries().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_queries() + Dnew, Inew = index.search(queries, k, numeric_type=numeric_type) cpu_index = faiss.index_gpu_to_cpu(index) + # cpu index always search in fp32 Dref, Iref = cpu_index.search(ds.get_queries(), k) evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) @@ -101,49 +78,55 @@ def do_interop(self, metric): faiss.serialize_index(cpu_index)) gpu_index = faiss.index_cpu_to_gpu(res, 0, deserialized_index) - Dnew2, Inew2 = gpu_index.search(ds.get_queries(), k) + Dnew2, Inew2 = gpu_index.search(queries, k, numeric_type=numeric_type) evaluation.check_ref_knn_with_draws(Dnew2, Inew2, Dnew, Inew, k) def test_interop_L2(self): - self.do_interop(faiss.METRIC_L2) + self.do_interop(faiss.METRIC_L2, faiss.Float32) def test_interop_IP(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT) + self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float32) + + def test_interop_L2_FP16(self): + self.do_interop(faiss.METRIC_L2, faiss.Float16) + + def test_interop_IP_FP16(self): + self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float16) + @unittest.skipIf( "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in") -class TestInteropFP16(unittest.TestCase): +class TestIDMapCagra(unittest.TestCase): - def do_interop(self, metric): + def do_IDMapCagra(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) + Dref, Iref = faiss.knn(ds.get_queries(), ds.get_database(), k, metric) res = faiss.StandardGpuResources() index = faiss.GpuIndexCagra(res, d, metric) - fp16_data = ds.get_database().astype(np.float16) - index.train(fp16_data, faiss.Float16) - fp16_queries = ds.get_queries().astype(np.float16) - Dnew, Inew = index.search(fp16_queries, k, numeric_type=faiss.Float16) + idMapIndex = faiss.IndexIDMap(index) + database = ds.get_database().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_database() + idMapIndex.train(database, numeric_type=numeric_type) + ids = np.array([i for i in range(10000)]) + idMapIndex.add_with_ids(database, ids, numeric_type=numeric_type) + queries = ds.get_queries().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_queries() + Dnew, Inew = idMapIndex.search(queries, k, numeric_type=numeric_type) - cpu_index = faiss.index_gpu_to_cpu(index) - Dref, Iref = cpu_index.search(ds.get_queries(), k) - evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) - deserialized_index = faiss.deserialize_index( - faiss.serialize_index(cpu_index)) + def test_IDMapCagra_L2(self): + self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float32) - gpu_index = faiss.index_cpu_to_gpu(res, 0, deserialized_index) - Dnew2, Inew2 = gpu_index.search(fp16_queries, k, numeric_type=faiss.Float16) + def test_IDMapCagra_IP(self): + self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float32) - evaluation.check_ref_knn_with_draws(Dnew2, Inew2, Dnew, Inew, k) - - def test_interop_L2(self): - self.do_interop(faiss.METRIC_L2) + def test_IDMapCagra_L2_FP16(self): + self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float16) - def test_interop_IP(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT) + def test_IDMapCagra_IP_FP16(self): + self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float16) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 51d8f570cb..848b84e190 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -42,6 +42,13 @@ def _check_dtype_uint8(codes): " uint8, but found %s" % ("codes", codes.dtype)) return np.ascontiguousarray(codes) +def _numeric_to_str(numeric_type): + if numeric_type == faiss.Float32: + return 'float32' + elif numeric_type == faiss.Float16: + return 'float16' + else: + raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ") def replace_method(the_class, name, replacement, ignore_missing=False): """ Replaces a method in a class with another version. The old method @@ -226,13 +233,13 @@ def replacement_add(self, x, numeric_type = faiss.Float32): n, d = x.shape assert d == self.d + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') + self.add_c(n, swig_ptr(x)) else: - x = np.ascontiguousarray(x, dtype='float16') - self.add_c(n, swig_ptr(x)) + self.add_c(n, swig_ptr(x), numeric_type) - def replacement_add_with_ids(self, x, ids): + def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): """Adds vectors with arbitrary ids to the index (not all indexes support this). The index must be trained before vectors can be added to it. Vector `i` is stored in `x[i]` and has id `ids[i]`. @@ -248,10 +255,14 @@ def replacement_add_with_ids(self, x, ids): """ n, d = x.shape assert d == self.d - x = np.ascontiguousarray(x, dtype='float32') - ids = np.ascontiguousarray(ids, dtype='int64') assert ids.shape == (n, ), 'not same nb of vectors as ids' - self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids)) + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) + ids = np.ascontiguousarray(ids, dtype='int64') + if numeric_type == faiss.Float32: + self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids)) + else: + self.add_with_ids_c(n, swig_ptr(x), numeric_type, swig_ptr(ids)) + def replacement_assign(self, x, k, labels=None): """Find the k nearest neighbors of the set of vectors x in the index. @@ -297,12 +308,11 @@ def replacement_train(self, x, numeric_type = faiss.Float32): """ n, d = x.shape assert d == self.d + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') self.train_c(n, swig_ptr(x)) else: - x = np.ascontiguousarray(x, dtype='float16') - self.train_c(n, swig_ptr(x), faiss.Float16) + self.train_c(n, swig_ptr(x), numeric_type) def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type = faiss.Float32): @@ -333,10 +343,7 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type """ n, d = x.shape - if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') - else: - x = np.ascontiguousarray(x, dtype='float16') + x = np.ascontiguousarray(x, _numeric_to_str(numeric_type)) assert d == self.d assert k > 0 @@ -354,7 +361,7 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type if numeric_type == faiss.Float32: self.search_c(n, swig_ptr(x), k, swig_ptr(D), swig_ptr(I), params) else: - self.search_c(n, swig_ptr(x), faiss.Float16, k, swig_ptr(D), swig_ptr(I), params) + self.search_c(n, swig_ptr(x), numeric_type, k, swig_ptr(D), swig_ptr(I), params) return D, I def replacement_search_and_reconstruct(self, x, k, *, params=None, D=None, I=None, R=None): @@ -893,7 +900,7 @@ def replacement_search(self, x, k, *, params=None): self.search_c(n, swig_ptr(x), k, swig_ptr(distances), swig_ptr(labels), - params=params) + params) return distances, labels def replacement_search_preassigned(self, x, k, Iq, Dq):