From 6ba70c3e5d9f7f41514692937a6090a058b3cadf Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 25 Jun 2025 01:01:50 +0000 Subject: [PATCH 01/16] IndexIDMap --- faiss/Index.h | 1 + faiss/IndexBinary.h | 53 ++++++++++++++++++++++++++++++++++ faiss/IndexIDMap.cpp | 48 ++++++++++++++++++++++++++---- faiss/IndexIDMap.h | 14 +++++++++ faiss/python/class_wrappers.py | 17 +++++++---- 5 files changed, 121 insertions(+), 12 deletions(-) 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/IndexBinary.h b/faiss/IndexBinary.h index e9801a7db4..0963aa2f0b 100644 --- a/faiss/IndexBinary.h +++ b/faiss/IndexBinary.h @@ -54,6 +54,15 @@ struct IndexBinary { * @param x training vecors, size n * d / 8 */ virtual void train(idx_t n, const uint8_t* x); + // This typed train function is a dummy to enable overriding of typed train + // function in the templated IndexIDMap struct. + 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 +70,15 @@ struct IndexBinary { * @param x input matrix, size n * d / 8 */ virtual void add(idx_t n, const uint8_t* x) = 0; + // This typed add function is a dummy to enable overriding of typed add + // function in the templated IndexIDMap struct. + 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 +88,20 @@ 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); + // This typed add_with_ids function is a dummy to enable overriding of typed + // add_with_ids function in the templated IndexIDMap struct. + 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 +119,27 @@ struct IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const = 0; + // This typed search function is a dummy to enable overriding of typed + // search function in the templated IndexIDMap struct. + 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/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index 6f4530544a..2b23f0f20e 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -12,6 +12,7 @@ #include #include #include +#include "faiss/Index.h" #include #include @@ -59,11 +60,19 @@ 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), NumericType::Float32); +} + template void IndexIDMapTemplate::reset() { index->reset(); @@ -74,14 +83,23 @@ 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), NumericType::Float32, xids); +} + template size_t IndexIDMapTemplate::sa_code_size() const { return index->sa_code_size(); @@ -123,7 +141,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 +166,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 +174,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), + NumericType::Float32, + k, + distances, + labels, + params); +} + template void IndexIDMapTemplate::range_search( idx_t n, diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h index dd3887ae76..42c8002dca 100644 --- a/faiss/IndexIDMap.h +++ b/faiss/IndexIDMap.h @@ -31,6 +31,11 @@ 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; @@ -42,8 +47,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; diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 51d8f570cb..ff2d2bbd73 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -230,9 +230,9 @@ def replacement_add(self, x, numeric_type = faiss.Float32): x = np.ascontiguousarray(x, dtype='float32') 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 +248,13 @@ def replacement_add_with_ids(self, x, ids): """ n, d = x.shape assert d == self.d - x = np.ascontiguousarray(x, dtype='float32') + if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='float32') + else: + x = np.ascontiguousarray(x, dtype='float16') 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)) + 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. @@ -295,6 +298,7 @@ def replacement_train(self, x, numeric_type = faiss.Float32): Query vectors, shape (n, d) where d is appropriate for the index. `dtype` must be float32. """ + print("we are here") n, d = x.shape assert d == self.d if numeric_type == faiss.Float32: @@ -302,7 +306,8 @@ def replacement_train(self, x, numeric_type = faiss.Float32): self.train_c(n, swig_ptr(x)) else: x = np.ascontiguousarray(x, dtype='float16') - self.train_c(n, swig_ptr(x), faiss.Float16) + print("we are here to train") + 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): @@ -350,7 +355,7 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type I = np.empty((n, k), dtype=np.int64) else: assert I.shape == (n, k) - + print("we are here in replacement search") if numeric_type == faiss.Float32: self.search_c(n, swig_ptr(x), k, swig_ptr(D), swig_ptr(I), params) else: From fdf4179cd0723fb16626de70c6c878e424a8e947 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 25 Jun 2025 17:41:54 +0000 Subject: [PATCH 02/16] rm numeric_type in python API + add tests --- faiss/gpu/test/test_cagra.py | 107 ++++++++++++++------------------- faiss/python/class_wrappers.py | 52 ++++++++-------- 2 files changed, 70 insertions(+), 89 deletions(-) diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 9c9297c888..11602bfe80 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, fp16): 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 fp16 else ds.get_database() + index.train(database) + queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() + Dnew, Inew = index.search(queries, k) 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, False) def test_compute_GT_IP(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT) + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, False) -@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, True) - 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, True) @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, fp16): 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 fp16 else ds.get_database() + index.train(database) + queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() + Dnew, Inew = index.search(queries, k) 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) 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, False) def test_interop_IP(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT) + self.do_interop(faiss.METRIC_INNER_PRODUCT, False) + + def test_interop_L2_FP16(self): + self.do_interop(faiss.METRIC_L2, True) + + def test_interop_IP_FP16(self): + self.do_interop(faiss.METRIC_INNER_PRODUCT, True) + @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, fp16): 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 fp16 else ds.get_database() + idMapIndex.train(database) + ids = [i for i in range(10000)] + idMapIndex.add_with_ids(database, ids) + queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() + Dnew, Inew = idMapIndex.search(queries, k) - 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, False) - 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, False) - 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, True) - 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, True) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index ff2d2bbd73..ad93b7ac98 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -42,6 +42,14 @@ def _check_dtype_uint8(codes): " uint8, but found %s" % ("codes", codes.dtype)) return np.ascontiguousarray(codes) +def _np_type_to_faiss_numeric(np_type): + if np_type == np.float32: + return faiss.Float32 + elif np_type == np.float16: + return faiss.Float16 + else: + raise TypeError("Input data must be ndarray of dtype " + " float32 or float16, but found %s" % (np_type)) def replace_method(the_class, name, replacement, ignore_missing=False): """ Replaces a method in a class with another version. The old method @@ -211,7 +219,7 @@ def replacement_build(self, x, graph): def handle_Index(the_class): - def replacement_add(self, x, numeric_type = faiss.Float32): + def replacement_add(self, x): """Adds vectors to the index. The index must be trained before vectors can be added to it. The vectors are implicitly numbered in sequence. When `n` vectors are @@ -226,13 +234,11 @@ def replacement_add(self, x, numeric_type = faiss.Float32): n, d = x.shape assert d == self.d - if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') - else: - x = np.ascontiguousarray(x, dtype='float16') + numeric_type = _np_type_to_faiss_numeric(x.dtype) + x = np.ascontiguousarray(x, dtype=x.dtype) self.add_c(n, swig_ptr(x), numeric_type) - def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): + def replacement_add_with_ids(self, x, ids): """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 +254,8 @@ def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): """ n, d = x.shape assert d == self.d - if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') - else: - x = np.ascontiguousarray(x, dtype='float16') + numeric_type = _np_type_to_faiss_numeric(x.dtype) + x = np.ascontiguousarray(x, dtype=x.dtype) 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), numeric_type, swig_ptr(ids)) @@ -288,7 +292,7 @@ def replacement_assign(self, x, k, labels=None): self.assign_c(n, swig_ptr(x), swig_ptr(labels), k) return labels - def replacement_train(self, x, numeric_type = faiss.Float32): + def replacement_train(self, x): """Trains the index on a representative set of vectors. The index must be trained before vectors can be added to it. @@ -298,19 +302,13 @@ def replacement_train(self, x, numeric_type = faiss.Float32): Query vectors, shape (n, d) where d is appropriate for the index. `dtype` must be float32. """ - print("we are here") n, d = x.shape assert d == self.d - 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') - print("we are here to train") - self.train_c(n, swig_ptr(x), numeric_type) - + numeric_type = _np_type_to_faiss_numeric(x.dtype) + x = np.ascontiguousarray(x, dtype=x.dtype) + 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): + def replacement_search(self, x, k, *, params=None, D=None, I=None): """Find the k nearest neighbors of the set of vectors x in the index. Parameters @@ -338,10 +336,9 @@ 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') + numeric_type = _np_type_to_faiss_numeric(x.dtype) + x = np.ascontiguousarray(x, dtype=x.dtype) + assert d == self.d assert k > 0 @@ -355,11 +352,12 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type I = np.empty((n, k), dtype=np.int64) else: assert I.shape == (n, k) - print("we are here in replacement search") + 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): From 95a681b2f1795fa070309b1c38f6d32373d22bdc Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 25 Jun 2025 20:26:36 +0000 Subject: [PATCH 03/16] fix errors --- faiss/python/class_wrappers.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index ad93b7ac98..ed4c94333a 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -236,7 +236,10 @@ def replacement_add(self, x): assert d == self.d numeric_type = _np_type_to_faiss_numeric(x.dtype) x = np.ascontiguousarray(x, dtype=x.dtype) - self.add_c(n, swig_ptr(x), numeric_type) + if numeric_type == faiss.Float32: + self.add_c(n, swig_ptr(x)) + else: + self.add_c(n, swig_ptr(x), numeric_type) def replacement_add_with_ids(self, x, ids): """Adds vectors with arbitrary ids to the index (not all indexes support this). @@ -258,7 +261,10 @@ def replacement_add_with_ids(self, x, ids): x = np.ascontiguousarray(x, dtype=x.dtype) 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), numeric_type, swig_ptr(ids)) + 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. @@ -306,7 +312,10 @@ def replacement_train(self, x): assert d == self.d numeric_type = _np_type_to_faiss_numeric(x.dtype) x = np.ascontiguousarray(x, dtype=x.dtype) - self.train_c(n, swig_ptr(x), numeric_type) + if numeric_type == faiss.Float32: + self.train_c(n, swig_ptr(x)) + else: + self.train_c(n, swig_ptr(x), numeric_type) def replacement_search(self, x, k, *, params=None, D=None, I=None): """Find the k nearest neighbors of the set of vectors x in the index. From 09543d6cc54ee90a96b3cecc508ef26378d9846c Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 25 Jun 2025 21:39:19 +0000 Subject: [PATCH 04/16] fix errors --- faiss/IndexIDMap.cpp | 23 +++++++++++++++++--- faiss/python/class_wrappers.py | 38 +++++++++++++++++++--------------- 2 files changed, 41 insertions(+), 20 deletions(-) diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index 2b23f0f20e..cfcff5fa46 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -34,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 *******************************************************/ @@ -70,7 +81,9 @@ template void IndexIDMapTemplate::train( idx_t n, const typename IndexT::component_t* x) { - train(n, static_cast(x), NumericType::Float32); + train(n, + static_cast(x), + component_t_to_numeric()); } template @@ -97,7 +110,11 @@ 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), NumericType::Float32, xids); + add_with_ids( + n, + static_cast(x), + component_t_to_numeric(), + xids); } template @@ -184,7 +201,7 @@ void IndexIDMapTemplate::search( const SearchParameters* params) const { search(n, static_cast(x), - NumericType::Float32, + component_t_to_numeric(), k, distances, labels, diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index ed4c94333a..e4045e7df3 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -43,13 +43,10 @@ def _check_dtype_uint8(codes): return np.ascontiguousarray(codes) def _np_type_to_faiss_numeric(np_type): - if np_type == np.float32: - return faiss.Float32 - elif np_type == np.float16: + if np_type == np.float16: return faiss.Float16 - else: - raise TypeError("Input data must be ndarray of dtype " - " float32 or float16, but found %s" % (np_type)) + else: # default + return faiss.Float32 def replace_method(the_class, name, replacement, ignore_missing=False): """ Replaces a method in a class with another version. The old method @@ -231,15 +228,16 @@ def replacement_add(self, x): Query vectors, shape (n, d) where d is appropriate for the index. `dtype` must be float32. """ - n, d = x.shape assert d == self.d numeric_type = _np_type_to_faiss_numeric(x.dtype) - x = np.ascontiguousarray(x, dtype=x.dtype) + if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='float32') self.add_c(n, swig_ptr(x)) - else: - self.add_c(n, swig_ptr(x), numeric_type) + else: # fp16 + x = np.ascontiguousarray(x, dtype='float16') + self.add_c(n, swig_ptr(x), faiss.Float16) def replacement_add_with_ids(self, x, ids): """Adds vectors with arbitrary ids to the index (not all indexes support this). @@ -258,13 +256,15 @@ def replacement_add_with_ids(self, x, ids): n, d = x.shape assert d == self.d numeric_type = _np_type_to_faiss_numeric(x.dtype) - x = np.ascontiguousarray(x, dtype=x.dtype) + ids = np.ascontiguousarray(ids, dtype='int64') assert ids.shape == (n, ), 'not same nb of vectors as ids' if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='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)) + else: # fp16 + x = np.ascontiguousarray(x, dtype='float16') + self.add_with_ids_c(n, swig_ptr(x), faiss.Float16, 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. @@ -311,11 +311,13 @@ def replacement_train(self, x): n, d = x.shape assert d == self.d numeric_type = _np_type_to_faiss_numeric(x.dtype) - x = np.ascontiguousarray(x, dtype=x.dtype) + if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='float32') self.train_c(n, swig_ptr(x)) else: - self.train_c(n, swig_ptr(x), numeric_type) + x = np.ascontiguousarray(x, dtype='float16') + self.train_c(n, swig_ptr(x), faiss.Float16) def replacement_search(self, x, k, *, params=None, D=None, I=None): """Find the k nearest neighbors of the set of vectors x in the index. @@ -346,7 +348,6 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None): n, d = x.shape numeric_type = _np_type_to_faiss_numeric(x.dtype) - x = np.ascontiguousarray(x, dtype=x.dtype) assert d == self.d @@ -363,9 +364,11 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None): assert I.shape == (n, k) if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='float32') self.search_c(n, swig_ptr(x), k, swig_ptr(D), swig_ptr(I), params) else: - self.search_c(n, swig_ptr(x), numeric_type, k, swig_ptr(D), swig_ptr(I), params) + x = np.ascontiguousarray(x, dtype='float16') + self.search_c(n, swig_ptr(x), faiss.Float16, k, swig_ptr(D), swig_ptr(I), params) return D, I @@ -860,6 +863,7 @@ def index_setstate(self, st): def handle_IndexBinary(the_class): def replacement_add(self, x): + print("inside uint") n, d = x.shape x = _check_dtype_uint8(x) assert d == self.code_size From d24b46d39a74f2700463520e92ba3f23da9891a4 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 9 Jul 2025 23:13:15 +0000 Subject: [PATCH 05/16] rever to numeric_type in python --- faiss/gpu/test/test_cagra.py | 60 +++++++++++++++++----------------- faiss/python/class_wrappers.py | 44 +++++++++---------------- 2 files changed, 46 insertions(+), 58 deletions(-) diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 11602bfe80..09f980f538 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, fp16): + def do_compute_GT(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) @@ -31,31 +31,31 @@ def do_compute_GT(self, metric, fp16): cagraIndexConfig.build_algo = faiss.graph_build_algo_IVF_PQ index = faiss.GpuIndexCagra(res, d, metric, cagraIndexConfig) - database = ds.get_database().astype(np.float16) if fp16 else ds.get_database() - index.train(database) - queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() - Dnew, Inew = index.search(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, False) + self.do_compute_GT(faiss.METRIC_L2, faiss.Float32) def test_compute_GT_IP(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, False) + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_compute_GT_L2_FP16(self): - self.do_compute_GT(faiss.METRIC_L2, True) + self.do_compute_GT(faiss.METRIC_L2, faiss.Float16) def test_compute_GT_IP_FP16(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, True) + 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, fp16): + def do_interop(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) @@ -63,10 +63,10 @@ def do_interop(self, metric, fp16): res = faiss.StandardGpuResources() index = faiss.GpuIndexCagra(res, d, metric) - database = ds.get_database().astype(np.float16) if fp16 else ds.get_database() - index.train(database) - queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() - Dnew, Inew = index.search(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 @@ -78,21 +78,21 @@ def do_interop(self, metric, fp16): faiss.serialize_index(cpu_index)) gpu_index = faiss.index_cpu_to_gpu(res, 0, deserialized_index) - Dnew2, Inew2 = gpu_index.search(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, False) + self.do_interop(faiss.METRIC_L2, faiss.Float32) def test_interop_IP(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT, False) + self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_interop_L2_FP16(self): - self.do_interop(faiss.METRIC_L2, True) + self.do_interop(faiss.METRIC_L2, faiss.Float16) def test_interop_IP_FP16(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT, True) + self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float16) @unittest.skipIf( @@ -100,7 +100,7 @@ def test_interop_IP_FP16(self): "only if cuVS is compiled in") class TestIDMapCagra(unittest.TestCase): - def do_IDMapCagra(self, metric, fp16): + def do_IDMapCagra(self, metric, numeric_type): d = 64 k = 12 ds = datasets.SyntheticDataset(d, 0, 10000, 100) @@ -110,23 +110,23 @@ def do_IDMapCagra(self, metric, fp16): index = faiss.GpuIndexCagra(res, d, metric) idMapIndex = faiss.IndexIDMap(index) - database = ds.get_database().astype(np.float16) if fp16 else ds.get_database() - idMapIndex.train(database) - ids = [i for i in range(10000)] - idMapIndex.add_with_ids(database, ids) - queries = ds.get_queries().astype(np.float16) if fp16 else ds.get_queries() - Dnew, Inew = idMapIndex.search(queries, k) + 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) evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) def test_IDMapCagra_L2(self): - self.do_IDMapCagra(faiss.METRIC_L2, False) + self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float32) def test_IDMapCagra_IP(self): - self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, False) + self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_IDMapCagra_L2_FP16(self): - self.do_IDMapCagra(faiss.METRIC_L2, True) + self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float16) def test_IDMapCagra_IP_FP16(self): - self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, True) + self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float16) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index e4045e7df3..eb7728c8d4 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -42,11 +42,6 @@ def _check_dtype_uint8(codes): " uint8, but found %s" % ("codes", codes.dtype)) return np.ascontiguousarray(codes) -def _np_type_to_faiss_numeric(np_type): - if np_type == np.float16: - return faiss.Float16 - else: # default - return faiss.Float32 def replace_method(the_class, name, replacement, ignore_missing=False): """ Replaces a method in a class with another version. The old method @@ -216,7 +211,7 @@ def replacement_build(self, x, graph): def handle_Index(the_class): - def replacement_add(self, x): + def replacement_add(self, x, numeric_type = faiss.Float32): """Adds vectors to the index. The index must be trained before vectors can be added to it. The vectors are implicitly numbered in sequence. When `n` vectors are @@ -228,18 +223,16 @@ def replacement_add(self, x): Query vectors, shape (n, d) where d is appropriate for the index. `dtype` must be float32. """ + n, d = x.shape assert d == self.d - numeric_type = _np_type_to_faiss_numeric(x.dtype) - if numeric_type == faiss.Float32: x = np.ascontiguousarray(x, dtype='float32') - self.add_c(n, swig_ptr(x)) - else: # fp16 + else: x = np.ascontiguousarray(x, dtype='float16') - self.add_c(n, swig_ptr(x), faiss.Float16) + self.add_c(n, swig_ptr(x)) - 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]`. @@ -255,14 +248,12 @@ def replacement_add_with_ids(self, x, ids): """ n, d = x.shape assert d == self.d - numeric_type = _np_type_to_faiss_numeric(x.dtype) - - ids = np.ascontiguousarray(ids, dtype='int64') assert ids.shape == (n, ), 'not same nb of vectors as ids' + if numeric_type == faiss.Float32: x = np.ascontiguousarray(x, dtype='float32') self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids)) - else: # fp16 + else: x = np.ascontiguousarray(x, dtype='float16') self.add_with_ids_c(n, swig_ptr(x), faiss.Float16, swig_ptr(ids)) @@ -298,7 +289,7 @@ def replacement_assign(self, x, k, labels=None): self.assign_c(n, swig_ptr(x), swig_ptr(labels), k) return labels - def replacement_train(self, x): + def replacement_train(self, x, numeric_type = faiss.Float32): """Trains the index on a representative set of vectors. The index must be trained before vectors can be added to it. @@ -310,16 +301,15 @@ def replacement_train(self, x): """ n, d = x.shape assert d == self.d - numeric_type = _np_type_to_faiss_numeric(x.dtype) - 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) + - def replacement_search(self, x, k, *, params=None, D=None, I=None): + def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type = faiss.Float32): """Find the k nearest neighbors of the set of vectors x in the index. Parameters @@ -347,8 +337,10 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None): """ n, d = x.shape - numeric_type = _np_type_to_faiss_numeric(x.dtype) - + if numeric_type == faiss.Float32: + x = np.ascontiguousarray(x, dtype='float32') + else: + x = np.ascontiguousarray(x, dtype='float16') assert d == self.d assert k > 0 @@ -362,14 +354,11 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None): I = np.empty((n, k), dtype=np.int64) else: assert I.shape == (n, k) - + if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') self.search_c(n, swig_ptr(x), k, swig_ptr(D), swig_ptr(I), params) else: - x = np.ascontiguousarray(x, dtype='float16') self.search_c(n, swig_ptr(x), faiss.Float16, 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): @@ -863,7 +852,6 @@ def index_setstate(self, st): def handle_IndexBinary(the_class): def replacement_add(self, x): - print("inside uint") n, d = x.shape x = _check_dtype_uint8(x) assert d == self.code_size @@ -1434,4 +1422,4 @@ def wrapper(*args, **kwargs): if len(args) > 3 and args[3] is not None: args[3] = faiss.PyCallbackShardingFunction(args[3]) return func(*args, **kwargs) - return wrapper + return wrapper \ No newline at end of file From 862c3a02f52c8b1f498776403d0b913b84c68cdc Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 9 Jul 2025 23:14:51 +0000 Subject: [PATCH 06/16] whitespaces --- faiss/gpu/test/test_cagra.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 09f980f538..9fd3e3b128 100644 --- a/faiss/gpu/test/test_cagra.py +++ b/faiss/gpu/test/test_cagra.py @@ -32,9 +32,9 @@ def do_compute_GT(self, metric, numeric_type): index = faiss.GpuIndexCagra(res, d, metric, cagraIndexConfig) database = ds.get_database().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_database() - index.train(database, numeric_type = numeric_type) + 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) + Dnew, Inew = index.search(queries, k, numeric_type=numeric_type) evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) @@ -42,13 +42,13 @@ def test_compute_GT_L2(self): self.do_compute_GT(faiss.METRIC_L2, faiss.Float32) def test_compute_GT_IP(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float32) + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_compute_GT_L2_FP16(self): - self.do_compute_GT(faiss.METRIC_L2, faiss.Float16) + self.do_compute_GT(faiss.METRIC_L2, faiss.Float16) def test_compute_GT_IP_FP16(self): - self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float16) + self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Float16) @unittest.skipIf( "CUVS" not in faiss.get_compile_options(), @@ -64,9 +64,9 @@ def do_interop(self, metric, numeric_type): index = faiss.GpuIndexCagra(res, d, metric) database = ds.get_database().astype(np.float16) if numeric_type == faiss.Float16 else ds.get_database() - index.train(database, numeric_type = numeric_type) + 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) + Dnew, Inew = index.search(queries, k, numeric_type=numeric_type) cpu_index = faiss.index_gpu_to_cpu(index) # cpu index always search in fp32 @@ -78,7 +78,7 @@ def do_interop(self, metric, numeric_type): faiss.serialize_index(cpu_index)) gpu_index = faiss.index_cpu_to_gpu(res, 0, deserialized_index) - Dnew2, Inew2 = gpu_index.search(queries, k, numeric_type = numeric_type) + Dnew2, Inew2 = gpu_index.search(queries, k, numeric_type=numeric_type) evaluation.check_ref_knn_with_draws(Dnew2, Inew2, Dnew, Inew, k) @@ -89,10 +89,10 @@ def test_interop_IP(self): self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_interop_L2_FP16(self): - self.do_interop(faiss.METRIC_L2, faiss.Float16) + self.do_interop(faiss.METRIC_L2, faiss.Float16) def test_interop_IP_FP16(self): - self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float16) + self.do_interop(faiss.METRIC_INNER_PRODUCT, faiss.Float16) @unittest.skipIf( @@ -111,11 +111,11 @@ def do_IDMapCagra(self, metric, numeric_type): index = faiss.GpuIndexCagra(res, d, metric) 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) + 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) + 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) + Dnew, Inew = idMapIndex.search(queries, k, numeric_type=numeric_type) evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k) @@ -126,7 +126,7 @@ def test_IDMapCagra_IP(self): self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float32) def test_IDMapCagra_L2_FP16(self): - self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float16) + self.do_IDMapCagra(faiss.METRIC_L2, faiss.Float16) def test_IDMapCagra_IP_FP16(self): - self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float16) + self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Float16) From b5c4dd411342431c7cef636f10303c00e903db9b Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 9 Jul 2025 23:25:16 +0000 Subject: [PATCH 07/16] fix add --- faiss/python/class_wrappers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index eb7728c8d4..12133a8bd4 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -228,9 +228,10 @@ def replacement_add(self, x, numeric_type = faiss.Float32): assert d == self.d 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), faiss.Float16) 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). From 610ce6a87f77cc74c6eec1758385f5f2bf8ae5cb Mon Sep 17 00:00:00 2001 From: jinsolp Date: Wed, 9 Jul 2025 23:26:13 +0000 Subject: [PATCH 08/16] newline --- faiss/python/class_wrappers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 12133a8bd4..bf3ba44472 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -1423,4 +1423,4 @@ def wrapper(*args, **kwargs): if len(args) > 3 and args[3] is not None: args[3] = faiss.PyCallbackShardingFunction(args[3]) return func(*args, **kwargs) - return wrapper \ No newline at end of file + return wrapper From 7a9982a6d2a46c2554b901eaeafdcb6e2fd7c0a5 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Thu, 10 Jul 2025 20:39:56 +0000 Subject: [PATCH 09/16] casting --- tests/test_contrib.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_contrib.py b/tests/test_contrib.py index a588362dbd..b25f0221dd 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -742,7 +742,7 @@ def do_test_ondisk_merge(self, shift_ids=False): index = faiss.read_index(tmpdir + "/trained.index") i0, i1 = int(bno * ds.nb / ns), int((bno + 1) * ds.nb / ns) if shift_ids: - index.add_with_ids(ds.xb[i0:i1], np.arange(0, ds.nb / ns)) + index.add_with_ids(ds.xb[i0:i1], np.arange(0, int(ds.nb / ns))) else: index.add_with_ids(ds.xb[i0:i1], np.arange(i0, i1)) faiss.write_index(index, tmpdir + "/block_%d.index" % bno) From a8d50193dcd6b811da352d4ea143c6654b36c91b Mon Sep 17 00:00:00 2001 From: jinsolp Date: Fri, 11 Jul 2025 03:34:18 +0000 Subject: [PATCH 10/16] fix Ex API and cleanup --- faiss/Index.h | 8 ++--- faiss/IndexBinary.h | 25 ++++++---------- faiss/IndexIDMap.cpp | 45 ++++++++++++++++++----------- faiss/IndexIDMap.h | 7 +++-- faiss/gpu/GpuIndex.cu | 35 +++++++++++----------- faiss/gpu/GpuIndex.h | 10 +++---- faiss/gpu/GpuIndexCagra.cu | 14 ++++----- faiss/gpu/GpuIndexCagra.h | 6 ++-- faiss/gpu/test/TestGpuIndexCagra.cu | 12 ++++---- faiss/python/class_wrappers.py | 43 ++++++++++----------------- 10 files changed, 100 insertions(+), 105 deletions(-) diff --git a/faiss/Index.h b/faiss/Index.h index 0af35cfce2..b1d1103521 100644 --- a/faiss/Index.h +++ b/faiss/Index.h @@ -127,7 +127,7 @@ struct Index { */ virtual void train(idx_t n, const float* x); - virtual void train(idx_t n, const void* x, NumericType numeric_type) { + virtual void trainEx(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::Float32) { train(n, static_cast(x)); } else { @@ -145,7 +145,7 @@ struct Index { */ virtual void add(idx_t n, const float* x) = 0; - virtual void add(idx_t n, const void* x, NumericType numeric_type) { + virtual void addEx(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::Float32) { add(n, static_cast(x)); } else { @@ -163,7 +163,7 @@ struct Index { * @param xids if non-null, ids to store for the vectors (size n) */ virtual void add_with_ids(idx_t n, const float* x, const idx_t* xids); - virtual void add_with_ids( + virtual void add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, @@ -194,7 +194,7 @@ struct Index { idx_t* labels, const SearchParameters* params = nullptr) const = 0; - virtual void search( + virtual void searchEx( idx_t n, const void* x, NumericType numeric_type, diff --git a/faiss/IndexBinary.h b/faiss/IndexBinary.h index 0963aa2f0b..c6c4e020a0 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,15 +55,13 @@ struct IndexBinary { * @param x training vecors, size n * d / 8 */ virtual void train(idx_t n, const uint8_t* x); - // This typed train function is a dummy to enable overriding of typed train - // function in the templated IndexIDMap struct. - virtual void train(idx_t n, const void* x, NumericType numeric_type) { + virtual void trainEx(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. * @@ -70,15 +69,13 @@ struct IndexBinary { * @param x input matrix, size n * d / 8 */ virtual void add(idx_t n, const uint8_t* x) = 0; - // This typed add function is a dummy to enable overriding of typed add - // function in the templated IndexIDMap struct. - virtual void add(idx_t n, const void* x, NumericType numeric_type) { + virtual void addEx(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. * @@ -88,9 +85,7 @@ 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); - // This typed add_with_ids function is a dummy to enable overriding of typed - // add_with_ids function in the templated IndexIDMap struct. - virtual void add_with_ids( + virtual void add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, @@ -101,7 +96,7 @@ struct IndexBinary { FAISS_THROW_MSG( "IndexBinary::add_with_ids: unsupported numeric type"); } - } + }; /** Query n vectors of dimension d to the index. * @@ -119,9 +114,7 @@ struct IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const = 0; - // This typed search function is a dummy to enable overriding of typed - // search function in the templated IndexIDMap struct. - virtual void search( + virtual void searchEx( idx_t n, const void* x, NumericType numeric_type, @@ -139,7 +132,7 @@ struct IndexBinary { } else { FAISS_THROW_MSG("IndexBinary::search: unsupported numeric type"); } - } + }; /** Query n vectors of dimension d to the index. * diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index cfcff5fa46..1ec429ca02 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -59,6 +59,16 @@ IndexIDMapTemplate::IndexIDMapTemplate(IndexT* index) : index(index) { sync_d(this); } +template +void IndexIDMapTemplate::addEx( + 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, @@ -69,11 +79,11 @@ void IndexIDMapTemplate::add( } template -void IndexIDMapTemplate::train( +void IndexIDMapTemplate::trainEx( idx_t n, const void* x, NumericType numeric_type) { - index->train(n, x, numeric_type); + index->trainEx(n, x, numeric_type); this->is_trained = index->is_trained; } @@ -81,9 +91,9 @@ template void IndexIDMapTemplate::train( idx_t n, const typename IndexT::component_t* x) { - train(n, - static_cast(x), - component_t_to_numeric()); + trainEx(n, + static_cast(x), + component_t_to_numeric()); } template @@ -94,12 +104,12 @@ void IndexIDMapTemplate::reset() { } template -void IndexIDMapTemplate::add_with_ids( +void IndexIDMapTemplate::add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, const idx_t* xids) { - index->add(n, x, numeric_type); + index->addEx(n, x, numeric_type); for (idx_t i = 0; i < n; i++) id_map.push_back(xids[i]); this->ntotal = index->ntotal; @@ -110,7 +120,7 @@ void IndexIDMapTemplate::add_with_ids( idx_t n, const typename IndexT::component_t* x, const idx_t* xids) { - add_with_ids( + add_with_idsEx( n, static_cast(x), component_t_to_numeric(), @@ -156,7 +166,7 @@ struct ScopedSelChange { } // namespace template -void IndexIDMapTemplate::search( +void IndexIDMapTemplate::searchEx( idx_t n, const void* x, NumericType numeric_type, @@ -183,7 +193,7 @@ void IndexIDMapTemplate::search( sel_change.set(params_non_const, &this_idtrans); } } - index->search(n, x, numeric_type, k, distances, labels, params); + index->searchEx(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++) { @@ -199,13 +209,14 @@ void IndexIDMapTemplate::search( 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); + searchEx( + n, + static_cast(x), + component_t_to_numeric(), + k, + distances, + labels, + params); } template diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h index 42c8002dca..5b47d0acec 100644 --- a/faiss/IndexIDMap.h +++ b/faiss/IndexIDMap.h @@ -31,7 +31,7 @@ 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( + void add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, @@ -39,6 +39,7 @@ struct IndexIDMapTemplate : IndexT { /// this will fail. Use add_with_ids void add(idx_t n, const component_t* x) override; + void addEx(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -47,7 +48,7 @@ struct IndexIDMapTemplate : IndexT { distance_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; - void search( + void searchEx( idx_t n, const void* x, NumericType numeric_type, @@ -57,7 +58,7 @@ struct IndexIDMapTemplate : IndexT { 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 trainEx(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; diff --git a/faiss/gpu/GpuIndex.cu b/faiss/gpu/GpuIndex.cu index 31c1bcddd1..d72b512239 100644 --- a/faiss/gpu/GpuIndex.cu +++ b/faiss/gpu/GpuIndex.cu @@ -110,16 +110,16 @@ size_t GpuIndex::getMinPagingSize() const { return minPagedSize_; } -void GpuIndex::add(idx_t n, const void* x, NumericType numeric_type) { - add_with_ids(n, x, numeric_type, nullptr); +void GpuIndex::addEx(idx_t n, const void* x, NumericType numeric_type) { + add_with_idsEx(n, x, numeric_type, nullptr); } void GpuIndex::add(idx_t n, const float* x) { // Pass to add_with_ids - add(n, x, NumericType::Float32); + addEx(n, x, NumericType::Float32); } -void GpuIndex::add_with_ids( +void GpuIndex::add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, @@ -147,7 +147,7 @@ void GpuIndex::add_with_ids( } void GpuIndex::add_with_ids(idx_t n, const float* x, const idx_t* ids) { - add_with_ids(n, static_cast(x), NumericType::Float32, ids); + add_with_idsEx(n, static_cast(x), NumericType::Float32, ids); } void GpuIndex::addPaged_( @@ -233,13 +233,13 @@ void GpuIndex::addPage_( stream, {n}); - addImpl_( + addImplEx_( n, static_cast(vecs.data()), numeric_type, ids ? indices.data() : nullptr); } else { - addImpl_( + addImplEx_( n, static_cast(vecs.data()), numeric_type, @@ -277,7 +277,7 @@ void GpuIndex::assign(idx_t n, const float* x, idx_t* labels, idx_t k) const { search(n, x, k, distances.data(), labels); } -void GpuIndex::search( +void GpuIndex::searchEx( idx_t n, const void* x, NumericType numeric_type, @@ -360,13 +360,14 @@ void GpuIndex::search( float* distances, idx_t* labels, const SearchParameters* params) const { - search(n, - static_cast(x), - NumericType::Float32, - k, - distances, - labels, - params); + searchEx( + n, + static_cast(x), + NumericType::Float32, + k, + distances, + labels, + params); } void GpuIndex::search_and_reconstruct( @@ -411,7 +412,7 @@ void GpuIndex::searchNonPaged_( stream, {n, this->d}); - searchImpl_( + searchImplEx_( n, static_cast(vecs.data()), numeric_type, @@ -600,7 +601,7 @@ void GpuIndex::searchFromCpuPaged_( auto outIndicesSlice = outIndices.narrowOutermost(cur3, numToProcess); - searchImpl_( + searchImplEx_( numToProcess, static_cast(bufGpus[cur3BufIndex]->data()), numeric_type, diff --git a/faiss/gpu/GpuIndex.h b/faiss/gpu/GpuIndex.h index bdeb362d31..e964032f87 100644 --- a/faiss/gpu/GpuIndex.h +++ b/faiss/gpu/GpuIndex.h @@ -77,13 +77,13 @@ class GpuIndex : public faiss::Index { /// as needed /// Handles paged adds if the add set is too large; calls addInternal_ void add(idx_t, const float* x) override; - void add(idx_t, const void* x, NumericType numeric_type) override; + void addEx(idx_t, const void* x, NumericType numeric_type) override; /// `x` and `ids` can be resident on the CPU or any GPU; copies are /// performed as needed /// Handles paged adds if the add set is too large; calls addInternal_ void add_with_ids(idx_t n, const float* x, const idx_t* ids) override; - void add_with_ids( + void add_with_idsEx( idx_t n, const void* x, NumericType numeric_type, @@ -103,7 +103,7 @@ class GpuIndex : public faiss::Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; - void search( + void searchEx( idx_t n, const void* x, NumericType numeric_type, @@ -165,7 +165,7 @@ class GpuIndex : public faiss::Index { /// All data is guaranteed to be resident on our device virtual void addImpl_(idx_t n, const float* x, const idx_t* ids) = 0; - virtual void addImpl_( + virtual void addImplEx_( idx_t n, const void* x, NumericType numeric_type, @@ -187,7 +187,7 @@ class GpuIndex : public faiss::Index { idx_t* labels, const SearchParameters* params) const = 0; - virtual void searchImpl_( + virtual void searchImplEx_( idx_t n, const void* x, NumericType numeric_type, diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu index 6bc4bc1cf5..0329f07540 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -42,7 +42,7 @@ GpuIndexCagra::GpuIndexCagra( this->is_trained = false; } -void GpuIndexCagra::train(idx_t n, const void* x, NumericType numeric_type) { +void GpuIndexCagra::trainEx(idx_t n, const void* x, NumericType numeric_type) { numeric_type_ = numeric_type; bool index_is_initialized = !std::holds_alternative(index_); @@ -133,15 +133,15 @@ void GpuIndexCagra::train(idx_t n, const void* x, NumericType numeric_type) { } void GpuIndexCagra::train(idx_t n, const float* x) { - train(n, static_cast(x), NumericType::Float32); + trainEx(n, static_cast(x), NumericType::Float32); } -void GpuIndexCagra::add(idx_t n, const void* x, NumericType numeric_type) { - train(n, x, numeric_type); +void GpuIndexCagra::addEx(idx_t n, const void* x, NumericType numeric_type) { + trainEx(n, x, numeric_type); } void GpuIndexCagra::add(idx_t n, const float* x) { - add(n, x, NumericType::Float32); + addEx(n, x, NumericType::Float32); } bool GpuIndexCagra::addImplRequiresIDs_() const { @@ -152,7 +152,7 @@ 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::searchImpl_( +void GpuIndexCagra::searchImplEx_( idx_t n, const void* x, NumericType numeric_type, @@ -240,7 +240,7 @@ void GpuIndexCagra::searchImpl_( float* distances, idx_t* labels, const SearchParameters* search_params) const { - searchImpl_( + searchImplEx_( n, static_cast(x), NumericType::Float32, diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index cf4a706e7d..0c5b6eeb4f 100644 --- a/faiss/gpu/GpuIndexCagra.h +++ b/faiss/gpu/GpuIndexCagra.h @@ -256,7 +256,7 @@ struct GpuIndexCagra : public GpuIndex { /// 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 float* x) override; - void add(idx_t n, const void* x, NumericType numeric_type) override; + void addEx(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 @@ -264,7 +264,7 @@ struct GpuIndexCagra : public GpuIndex { /// of vectors (without IDs) to the index. There is no external quantizer to /// be trained here. void train(idx_t n, const float* x) override; - void train(idx_t n, const void* x, NumericType numeric_type) override; + void trainEx(idx_t n, const void* x, NumericType numeric_type) override; /// Initialize ourselves from the given CPU index; will overwrite /// all data in ourselves @@ -294,7 +294,7 @@ struct GpuIndexCagra : public GpuIndex { float* distances, idx_t* labels, const SearchParameters* search_params) const override; - void searchImpl_( + void searchImplEx_( idx_t n, const void* x, NumericType numeric_type, diff --git a/faiss/gpu/test/TestGpuIndexCagra.cu b/faiss/gpu/test/TestGpuIndexCagra.cu index 31dcfc03c9..5693760d19 100644 --- a/faiss/gpu/test/TestGpuIndexCagra.cu +++ b/faiss/gpu/test/TestGpuIndexCagra.cu @@ -234,7 +234,7 @@ void queryTestFP16(faiss::MetricType metric, double expected_recall) { trainVecs_half[i] = __float2half(trainVecs[i]); } - gpuIndex.train( + gpuIndex.trainEx( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -272,7 +272,7 @@ void queryTestFP16(faiss::MetricType metric, double expected_recall) { for (size_t i = 0; i < queryVecs.size(); ++i) { queryVecs_half[i] = __float2half(queryVecs[i]); } - gpuIndex.search( + gpuIndex.searchEx( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, @@ -527,7 +527,7 @@ void copyToTestFP16( } faiss::gpu::GpuIndexCagra gpuIndex(&res, opt.dim, metric, config); - gpuIndex.train( + gpuIndex.trainEx( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -803,7 +803,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { trainVecs_half[i] = __float2half(trainVecs[i]); } - gpuIndex.train( + gpuIndex.trainEx( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -829,7 +829,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); faiss::gpu::DeviceTensor copyTestIndices( gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); - copiedGpuIndex.search( + copiedGpuIndex.searchEx( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, @@ -841,7 +841,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); faiss::gpu::DeviceTensor testIndices( gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); - gpuIndex.search( + gpuIndex.searchEx( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index bf3ba44472..4b0e8bdf84 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,12 +233,8 @@ def replacement_add(self, x, numeric_type = faiss.Float32): n, d = x.shape assert d == self.d - 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), faiss.Float16) + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) + self.addEx(n, swig_ptr(x), numeric_type) 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). @@ -250,13 +253,9 @@ def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): n, d = x.shape assert d == self.d assert ids.shape == (n, ), 'not same nb of vectors as ids' - - if numeric_type == faiss.Float32: - x = np.ascontiguousarray(x, dtype='float32') - self.add_with_ids_c(n, swig_ptr(x), swig_ptr(ids)) - else: - x = np.ascontiguousarray(x, dtype='float16') - self.add_with_ids_c(n, swig_ptr(x), faiss.Float16, swig_ptr(ids)) + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) + self.add_with_idsEx(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. @@ -302,12 +301,8 @@ def replacement_train(self, x, numeric_type = faiss.Float32): """ n, d = x.shape assert d == self.d - 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) + x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) + self.trainEx(n, swig_ptr(x), numeric_type) def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type = faiss.Float32): @@ -338,10 +333,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 @@ -356,10 +348,7 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type else: assert I.shape == (n, k) - 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.searchEx(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): From 76d59a88901b1b72fc31701ae6d040bb2b7a7bf2 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Fri, 11 Jul 2025 04:26:24 +0000 Subject: [PATCH 11/16] torch utils use extended API --- contrib/torch_utils.py | 55 +++++++++++++++++++++++++++++------------- 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py index 797c02656c..e5c6312d80 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.addEx(n, x_ptr, numeric_type) else: # CPU torch - self.add_c(n, x_ptr) + self.addEx(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_idsEx(n, x_ptr, numeric_type, ids_ptr) else: # CPU torch - self.add_with_ids_c(n, x_ptr, ids_ptr) + self.add_with_idsEx(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.trainEx(n, x_ptr, numeric_type) else: # CPU torch - self.train_c(n, x_ptr) + self.trainEx(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.searchEx(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.searchEx(n, x_ptr, numeric_type, k, D_ptr, I_ptr) return D, I From 0f4926c81c5c4e3c0bc158586d181e7f44018c83 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Fri, 11 Jul 2025 06:04:47 +0000 Subject: [PATCH 12/16] ex addwithids for indexmap2 --- faiss/IndexIDMap.cpp | 19 ++++++++++++++++--- faiss/IndexIDMap.h | 5 +++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index 1ec429ca02..87204d88a0 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -300,17 +300,30 @@ IndexIDMap2Template::IndexIDMap2Template(IndexT* index) : IndexIDMapTemplate(index) {} template -void IndexIDMap2Template::add_with_ids( +void IndexIDMap2Template::add_with_idsEx( 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_idsEx(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_idsEx( + 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 5b47d0acec..6286186ccc 100644 --- a/faiss/IndexIDMap.h +++ b/faiss/IndexIDMap.h @@ -104,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_idsEx( + idx_t n, + const void* x, + NumericType numeric_type, + const idx_t* xids) override; size_t remove_ids(const IDSelector& sel) override; From 8159855938efb1d34da937060ffc7b21aca0c40c Mon Sep 17 00:00:00 2001 From: jinsolp Date: Fri, 11 Jul 2025 07:31:50 +0000 Subject: [PATCH 13/16] bring back removed line --- faiss/python/class_wrappers.py | 1 + tests/test_contrib.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 4b0e8bdf84..100816f94f 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -254,6 +254,7 @@ def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): assert d == self.d assert ids.shape == (n, ), 'not same nb of vectors as ids' x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) + ids = np.ascontiguousarray(ids, dtype='int64') self.add_with_idsEx(n, swig_ptr(x), numeric_type, swig_ptr(ids)) diff --git a/tests/test_contrib.py b/tests/test_contrib.py index b25f0221dd..a588362dbd 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -742,7 +742,7 @@ def do_test_ondisk_merge(self, shift_ids=False): index = faiss.read_index(tmpdir + "/trained.index") i0, i1 = int(bno * ds.nb / ns), int((bno + 1) * ds.nb / ns) if shift_ids: - index.add_with_ids(ds.xb[i0:i1], np.arange(0, int(ds.nb / ns))) + index.add_with_ids(ds.xb[i0:i1], np.arange(0, ds.nb / ns)) else: index.add_with_ids(ds.xb[i0:i1], np.arange(i0, i1)) faiss.write_index(index, tmpdir + "/block_%d.index" % bno) From 74f69e41ba0150c9abdae5f5514a2b739005c93b Mon Sep 17 00:00:00 2001 From: jinsolp Date: Mon, 21 Jul 2025 21:53:10 +0000 Subject: [PATCH 14/16] remove Ex api --- contrib/torch_utils.py | 16 +++---- faiss/Index.h | 8 ++-- faiss/IndexAdditiveQuantizer.cpp | 25 +++++++++++ faiss/IndexAdditiveQuantizer.h | 11 +++++ faiss/IndexBinary.h | 8 ++-- faiss/IndexBinaryFlat.cpp | 15 +++++++ faiss/IndexBinaryFlat.h | 9 ++++ faiss/IndexBinaryHNSW.cpp | 19 ++++++++ faiss/IndexBinaryHNSW.h | 10 +++++ faiss/IndexFastScan.cpp | 15 +++++++ faiss/IndexFastScan.h | 9 ++++ faiss/IndexFlat.cpp | 26 +++++++++++ faiss/IndexFlat.h | 17 ++++++++ faiss/IndexFlatCodes.cpp | 15 +++++++ faiss/IndexFlatCodes.h | 9 ++++ faiss/IndexHNSW.cpp | 49 +++++++++++++++++++++ faiss/IndexHNSW.h | 28 ++++++++++++ faiss/IndexIDMap.cpp | 43 +++++++++--------- faiss/IndexIDMap.h | 10 ++--- faiss/IndexIVF.cpp | 27 ++++++++++++ faiss/IndexIVF.h | 15 +++++++ faiss/IndexIVFFlat.cpp | 15 +++++++ faiss/IndexIVFFlat.h | 6 +++ faiss/IndexNNDescent.cpp | 19 ++++++++ faiss/IndexNNDescent.h | 11 ++++- faiss/IndexPQ.cpp | 58 +++++++++++++++++++++++++ faiss/IndexPQ.h | 29 +++++++++++++ faiss/IndexScalarQuantizer.cpp | 18 ++++++++ faiss/IndexScalarQuantizer.h | 9 ++++ faiss/gpu/GpuIndex.cu | 27 ++++++------ faiss/gpu/GpuIndex.h | 6 +-- faiss/gpu/GpuIndexBinaryCagra.cu | 25 +++++++++++ faiss/gpu/GpuIndexBinaryCagra.h | 10 +++++ faiss/gpu/GpuIndexBinaryFlat.cu | 15 +++++++ faiss/gpu/GpuIndexBinaryFlat.h | 9 ++++ faiss/gpu/GpuIndexCagra.cu | 10 ++--- faiss/gpu/GpuIndexCagra.h | 4 +- faiss/gpu/GpuIndexFlat.cu | 8 ++++ faiss/gpu/GpuIndexFlat.h | 2 + faiss/gpu/GpuIndexIVFFlat.cu | 4 ++ faiss/gpu/GpuIndexIVFFlat.h | 1 + faiss/gpu/GpuIndexIVFPQ.cu | 4 ++ faiss/gpu/GpuIndexIVFPQ.h | 1 + faiss/gpu/GpuIndexIVFScalarQuantizer.cu | 7 +++ faiss/gpu/GpuIndexIVFScalarQuantizer.h | 1 + faiss/gpu/test/TestGpuIndexCagra.cu | 12 ++--- faiss/python/class_wrappers.py | 22 +++++++--- 47 files changed, 638 insertions(+), 79 deletions(-) diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py index e5c6312d80..0056ba01ec 100644 --- a/contrib/torch_utils.py +++ b/contrib/torch_utils.py @@ -152,10 +152,10 @@ def torch_replacement_add(self, x, numeric_type = faiss.Float32): # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.addEx(n, x_ptr, numeric_type) + self.add_c(n, x_ptr, numeric_type) else: # CPU torch - self.addEx(n, x_ptr, numeric_type) + self.add_c(n, x_ptr, numeric_type) def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): if type(x) is np.ndarray: @@ -181,10 +181,10 @@ def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.add_with_idsEx(n, x_ptr, numeric_type, ids_ptr) + self.add_with_ids_c(n, x_ptr, numeric_type, ids_ptr) else: # CPU torch - self.add_with_idsEx(n, x_ptr, numeric_type, 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: @@ -235,10 +235,10 @@ def torch_replacement_train(self, x, numeric_type = faiss.Float32): # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.trainEx(n, x_ptr, numeric_type) + self.train_c(n, x_ptr, numeric_type) else: # CPU torch - self.trainEx(n, x_ptr, numeric_type) + self.train_c(n, x_ptr, numeric_type) def search_methods_common(x, k, D, I, numeric_type=faiss.Float32): n, d = x.shape @@ -281,10 +281,10 @@ def torch_replacement_search(self, x, k, D=None, I=None, numeric_type=faiss.Floa # On the GPU, use proper stream ordering with using_stream(self.getResources()): - self.searchEx(n, x_ptr, numeric_type, k, D_ptr, I_ptr) + self.search_c(n, x_ptr, numeric_type, k, D_ptr, I_ptr) else: # CPU torch - self.searchEx(n, x_ptr, numeric_type, 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 b1d1103521..0af35cfce2 100644 --- a/faiss/Index.h +++ b/faiss/Index.h @@ -127,7 +127,7 @@ struct Index { */ virtual void train(idx_t n, const float* x); - virtual void trainEx(idx_t n, const void* x, NumericType numeric_type) { + virtual void train(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::Float32) { train(n, static_cast(x)); } else { @@ -145,7 +145,7 @@ struct Index { */ virtual void add(idx_t n, const float* x) = 0; - virtual void addEx(idx_t n, const void* x, NumericType numeric_type) { + virtual void add(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::Float32) { add(n, static_cast(x)); } else { @@ -163,7 +163,7 @@ struct Index { * @param xids if non-null, ids to store for the vectors (size n) */ virtual void add_with_ids(idx_t n, const float* x, const idx_t* xids); - virtual void add_with_idsEx( + virtual void add_with_ids( idx_t n, const void* x, NumericType numeric_type, @@ -194,7 +194,7 @@ struct Index { idx_t* labels, const SearchParameters* params = nullptr) const = 0; - virtual void searchEx( + virtual void search( idx_t n, const void* x, 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 c6c4e020a0..d1725921f7 100644 --- a/faiss/IndexBinary.h +++ b/faiss/IndexBinary.h @@ -55,7 +55,7 @@ struct IndexBinary { * @param x training vecors, size n * d / 8 */ virtual void train(idx_t n, const uint8_t* x); - virtual void trainEx(idx_t n, const void* x, NumericType numeric_type) { + virtual void train(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::UInt8) { train(n, static_cast(x)); } else { @@ -69,7 +69,7 @@ struct IndexBinary { * @param x input matrix, size n * d / 8 */ virtual void add(idx_t n, const uint8_t* x) = 0; - virtual void addEx(idx_t n, const void* x, NumericType numeric_type) { + virtual void add(idx_t n, const void* x, NumericType numeric_type) { if (numeric_type == NumericType::UInt8) { add(n, static_cast(x)); } else { @@ -85,7 +85,7 @@ 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_idsEx( + virtual void add_with_ids( idx_t n, const void* x, NumericType numeric_type, @@ -114,7 +114,7 @@ struct IndexBinary { int32_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const = 0; - virtual void searchEx( + virtual void search( idx_t n, const void* x, NumericType numeric_type, 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 87204d88a0..753254da04 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -60,7 +60,7 @@ IndexIDMapTemplate::IndexIDMapTemplate(IndexT* index) : index(index) { } template -void IndexIDMapTemplate::addEx( +void IndexIDMapTemplate::add( idx_t, const void*, NumericType numeric_type) { @@ -79,11 +79,11 @@ void IndexIDMapTemplate::add( } template -void IndexIDMapTemplate::trainEx( +void IndexIDMapTemplate::train( idx_t n, const void* x, NumericType numeric_type) { - index->trainEx(n, x, numeric_type); + index->train(n, x, numeric_type); this->is_trained = index->is_trained; } @@ -91,9 +91,9 @@ template void IndexIDMapTemplate::train( idx_t n, const typename IndexT::component_t* x) { - trainEx(n, - static_cast(x), - component_t_to_numeric()); + train(n, + static_cast(x), + component_t_to_numeric()); } template @@ -104,12 +104,12 @@ void IndexIDMapTemplate::reset() { } template -void IndexIDMapTemplate::add_with_idsEx( +void IndexIDMapTemplate::add_with_ids( idx_t n, const void* x, NumericType numeric_type, const idx_t* xids) { - index->addEx(n, x, numeric_type); + index->add(n, x, numeric_type); for (idx_t i = 0; i < n; i++) id_map.push_back(xids[i]); this->ntotal = index->ntotal; @@ -120,7 +120,7 @@ void IndexIDMapTemplate::add_with_ids( idx_t n, const typename IndexT::component_t* x, const idx_t* xids) { - add_with_idsEx( + add_with_ids( n, static_cast(x), component_t_to_numeric(), @@ -166,7 +166,7 @@ struct ScopedSelChange { } // namespace template -void IndexIDMapTemplate::searchEx( +void IndexIDMapTemplate::search( idx_t n, const void* x, NumericType numeric_type, @@ -193,7 +193,7 @@ void IndexIDMapTemplate::searchEx( sel_change.set(params_non_const, &this_idtrans); } } - index->searchEx(n, x, numeric_type, 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++) { @@ -209,14 +209,13 @@ void IndexIDMapTemplate::search( typename IndexT::distance_t* distances, idx_t* labels, const SearchParameters* params) const { - searchEx( - n, - static_cast(x), - component_t_to_numeric(), - k, - distances, - labels, - params); + search(n, + static_cast(x), + component_t_to_numeric(), + k, + distances, + labels, + params); } template @@ -300,13 +299,13 @@ IndexIDMap2Template::IndexIDMap2Template(IndexT* index) : IndexIDMapTemplate(index) {} template -void IndexIDMap2Template::add_with_idsEx( +void IndexIDMap2Template::add_with_ids( idx_t n, const void* x, NumericType numeric_type, const idx_t* xids) { size_t prev_ntotal = this->ntotal; - IndexIDMapTemplate::add_with_idsEx(n, x, numeric_type, 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; } @@ -317,7 +316,7 @@ void IndexIDMap2Template::add_with_ids( idx_t n, const typename IndexT::component_t* x, const idx_t* xids) { - add_with_idsEx( + add_with_ids( n, static_cast(x), component_t_to_numeric(), diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h index 6286186ccc..cfe43377ee 100644 --- a/faiss/IndexIDMap.h +++ b/faiss/IndexIDMap.h @@ -31,7 +31,7 @@ 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_idsEx( + void add_with_ids( idx_t n, const void* x, NumericType numeric_type, @@ -39,7 +39,7 @@ struct IndexIDMapTemplate : IndexT { /// this will fail. Use add_with_ids void add(idx_t n, const component_t* x) override; - void addEx(idx_t n, const void* x, NumericType numeric_type) override; + void add(idx_t n, const void* x, NumericType numeric_type) override; void search( idx_t n, @@ -48,7 +48,7 @@ struct IndexIDMapTemplate : IndexT { distance_t* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; - void searchEx( + void search( idx_t n, const void* x, NumericType numeric_type, @@ -58,7 +58,7 @@ struct IndexIDMapTemplate : IndexT { const SearchParameters* params = nullptr) const override; void train(idx_t n, const component_t* x) override; - void trainEx(idx_t n, const void* x, NumericType numeric_type) override; + void train(idx_t n, const void* x, NumericType numeric_type) override; void reset() override; @@ -104,7 +104,7 @@ struct IndexIDMap2Template : IndexIDMapTemplate { void add_with_ids(idx_t n, const component_t* x, const idx_t* xids) override; - void add_with_idsEx( + void add_with_ids( idx_t n, const void* x, NumericType numeric_type, 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/GpuIndex.cu b/faiss/gpu/GpuIndex.cu index d72b512239..4ab3f068c8 100644 --- a/faiss/gpu/GpuIndex.cu +++ b/faiss/gpu/GpuIndex.cu @@ -110,16 +110,16 @@ size_t GpuIndex::getMinPagingSize() const { return minPagedSize_; } -void GpuIndex::addEx(idx_t n, const void* x, NumericType numeric_type) { - add_with_idsEx(n, x, numeric_type, nullptr); +void GpuIndex::add(idx_t n, const void* x, NumericType numeric_type) { + add_with_ids(n, x, numeric_type, nullptr); } void GpuIndex::add(idx_t n, const float* x) { // Pass to add_with_ids - addEx(n, x, NumericType::Float32); + add(n, x, NumericType::Float32); } -void GpuIndex::add_with_idsEx( +void GpuIndex::add_with_ids( idx_t n, const void* x, NumericType numeric_type, @@ -147,7 +147,7 @@ void GpuIndex::add_with_idsEx( } void GpuIndex::add_with_ids(idx_t n, const float* x, const idx_t* ids) { - add_with_idsEx(n, static_cast(x), NumericType::Float32, ids); + add_with_ids(n, static_cast(x), NumericType::Float32, ids); } void GpuIndex::addPaged_( @@ -277,7 +277,7 @@ void GpuIndex::assign(idx_t n, const float* x, idx_t* labels, idx_t k) const { search(n, x, k, distances.data(), labels); } -void GpuIndex::searchEx( +void GpuIndex::search( idx_t n, const void* x, NumericType numeric_type, @@ -360,14 +360,13 @@ void GpuIndex::search( float* distances, idx_t* labels, const SearchParameters* params) const { - searchEx( - n, - static_cast(x), - NumericType::Float32, - k, - distances, - labels, - params); + search(n, + static_cast(x), + NumericType::Float32, + k, + distances, + labels, + params); } void GpuIndex::search_and_reconstruct( diff --git a/faiss/gpu/GpuIndex.h b/faiss/gpu/GpuIndex.h index e964032f87..eb1f7b2d01 100644 --- a/faiss/gpu/GpuIndex.h +++ b/faiss/gpu/GpuIndex.h @@ -77,13 +77,13 @@ class GpuIndex : public faiss::Index { /// as needed /// Handles paged adds if the add set is too large; calls addInternal_ void add(idx_t, const float* x) override; - void addEx(idx_t, const void* x, NumericType numeric_type) override; + void add(idx_t, const void* x, NumericType numeric_type) override; /// `x` and `ids` can be resident on the CPU or any GPU; copies are /// performed as needed /// Handles paged adds if the add set is too large; calls addInternal_ void add_with_ids(idx_t n, const float* x, const idx_t* ids) override; - void add_with_idsEx( + void add_with_ids( idx_t n, const void* x, NumericType numeric_type, @@ -103,7 +103,7 @@ class GpuIndex : public faiss::Index { float* distances, idx_t* labels, const SearchParameters* params = nullptr) const override; - void searchEx( + void search( idx_t n, const void* x, NumericType numeric_type, 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 0329f07540..6f129fe1d3 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -42,7 +42,7 @@ GpuIndexCagra::GpuIndexCagra( this->is_trained = false; } -void GpuIndexCagra::trainEx(idx_t n, const void* x, NumericType numeric_type) { +void GpuIndexCagra::train(idx_t n, const void* x, NumericType numeric_type) { numeric_type_ = numeric_type; bool index_is_initialized = !std::holds_alternative(index_); @@ -133,15 +133,15 @@ void GpuIndexCagra::trainEx(idx_t n, const void* x, NumericType numeric_type) { } void GpuIndexCagra::train(idx_t n, const float* x) { - trainEx(n, static_cast(x), NumericType::Float32); + train(n, static_cast(x), NumericType::Float32); } -void GpuIndexCagra::addEx(idx_t n, const void* x, NumericType numeric_type) { - trainEx(n, x, numeric_type); +void GpuIndexCagra::add(idx_t n, const void* x, NumericType numeric_type) { + train(n, x, numeric_type); } void GpuIndexCagra::add(idx_t n, const float* x) { - addEx(n, x, NumericType::Float32); + add(n, x, NumericType::Float32); } bool GpuIndexCagra::addImplRequiresIDs_() const { diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index 0c5b6eeb4f..1d60e3a72b 100644 --- a/faiss/gpu/GpuIndexCagra.h +++ b/faiss/gpu/GpuIndexCagra.h @@ -256,7 +256,7 @@ struct GpuIndexCagra : public GpuIndex { /// 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 float* x) override; - void addEx(idx_t n, const void* x, NumericType numeric_type) 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 @@ -264,7 +264,7 @@ struct GpuIndexCagra : public GpuIndex { /// of vectors (without IDs) to the index. There is no external quantizer to /// be trained here. void train(idx_t n, const float* x) override; - void trainEx(idx_t n, const void* x, NumericType numeric_type) 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 diff --git a/faiss/gpu/GpuIndexFlat.cu b/faiss/gpu/GpuIndexFlat.cu index eb87e082e9..78ed2410ad 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) { + Index::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) { + Index::add(n, x, numeric_type); +} + bool GpuIndexFlat::addImplRequiresIDs_() const { return false; } diff --git a/faiss/gpu/GpuIndexFlat.h b/faiss/gpu/GpuIndexFlat.h index ee4c14466e..62dc049ded 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 diff --git a/faiss/gpu/GpuIndexIVFFlat.cu b/faiss/gpu/GpuIndexIVFFlat.cu index 1266b992cf..4a85795ac1 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) { + Index::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..5cd1ec1eb3 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) { + Index::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..d5b0e81e68 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) { + Index::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/TestGpuIndexCagra.cu b/faiss/gpu/test/TestGpuIndexCagra.cu index 5693760d19..31dcfc03c9 100644 --- a/faiss/gpu/test/TestGpuIndexCagra.cu +++ b/faiss/gpu/test/TestGpuIndexCagra.cu @@ -234,7 +234,7 @@ void queryTestFP16(faiss::MetricType metric, double expected_recall) { trainVecs_half[i] = __float2half(trainVecs[i]); } - gpuIndex.trainEx( + gpuIndex.train( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -272,7 +272,7 @@ void queryTestFP16(faiss::MetricType metric, double expected_recall) { for (size_t i = 0; i < queryVecs.size(); ++i) { queryVecs_half[i] = __float2half(queryVecs[i]); } - gpuIndex.searchEx( + gpuIndex.search( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, @@ -527,7 +527,7 @@ void copyToTestFP16( } faiss::gpu::GpuIndexCagra gpuIndex(&res, opt.dim, metric, config); - gpuIndex.trainEx( + gpuIndex.train( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -803,7 +803,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { trainVecs_half[i] = __float2half(trainVecs[i]); } - gpuIndex.trainEx( + gpuIndex.train( opt.numTrain, static_cast(trainVecs_half.data()), faiss::NumericType::Float16); @@ -829,7 +829,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); faiss::gpu::DeviceTensor copyTestIndices( gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); - copiedGpuIndex.searchEx( + copiedGpuIndex.search( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, @@ -841,7 +841,7 @@ void copyFromTestFP16(faiss::MetricType metric, double expected_recall) { gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); faiss::gpu::DeviceTensor testIndices( gpuRes.get(), devAlloc, {opt.numQuery, opt.k}); - gpuIndex.searchEx( + gpuIndex.search( opt.numQuery, queryVecs_half.data(), faiss::NumericType::Float16, diff --git a/faiss/python/class_wrappers.py b/faiss/python/class_wrappers.py index 100816f94f..848b84e190 100644 --- a/faiss/python/class_wrappers.py +++ b/faiss/python/class_wrappers.py @@ -234,7 +234,10 @@ 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)) - self.addEx(n, swig_ptr(x), numeric_type) + if numeric_type == faiss.Float32: + self.add_c(n, swig_ptr(x)) + else: + self.add_c(n, swig_ptr(x), numeric_type) 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). @@ -255,7 +258,10 @@ def replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32): assert ids.shape == (n, ), 'not same nb of vectors as ids' x = np.ascontiguousarray(x, dtype=_numeric_to_str(numeric_type)) ids = np.ascontiguousarray(ids, dtype='int64') - self.add_with_idsEx(n, swig_ptr(x), numeric_type, swig_ptr(ids)) + 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): @@ -303,7 +309,10 @@ 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)) - self.trainEx(n, swig_ptr(x), numeric_type) + if numeric_type == faiss.Float32: + self.train_c(n, swig_ptr(x)) + else: + 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): @@ -349,7 +358,10 @@ def replacement_search(self, x, k, *, params=None, D=None, I=None, numeric_type else: assert I.shape == (n, k) - self.searchEx(n, swig_ptr(x), numeric_type, k, swig_ptr(D), swig_ptr(I), params) + 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), 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): @@ -888,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): From 96842cd2e1a7ce0d4d9cabfebeaf847ebdd7b1e9 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Mon, 21 Jul 2025 22:52:37 +0000 Subject: [PATCH 15/16] rm GpuIndex Ex api --- faiss/gpu/GpuIndex.cu | 8 ++++---- faiss/gpu/GpuIndex.h | 4 ++-- faiss/gpu/GpuIndexCagra.cu | 12 ++++++++++-- faiss/gpu/GpuIndexCagra.h | 7 ++++++- faiss/gpu/GpuIndexFlat.cu | 19 +++++++++++++++++++ faiss/gpu/GpuIndexFlat.h | 13 +++++++++++++ faiss/gpu/GpuIndexIVF.cu | 20 ++++++++++++++++++++ faiss/gpu/GpuIndexIVF.h | 14 ++++++++++++++ 8 files changed, 88 insertions(+), 9 deletions(-) diff --git a/faiss/gpu/GpuIndex.cu b/faiss/gpu/GpuIndex.cu index 4ab3f068c8..31c1bcddd1 100644 --- a/faiss/gpu/GpuIndex.cu +++ b/faiss/gpu/GpuIndex.cu @@ -233,13 +233,13 @@ void GpuIndex::addPage_( stream, {n}); - addImplEx_( + addImpl_( n, static_cast(vecs.data()), numeric_type, ids ? indices.data() : nullptr); } else { - addImplEx_( + addImpl_( n, static_cast(vecs.data()), numeric_type, @@ -411,7 +411,7 @@ void GpuIndex::searchNonPaged_( stream, {n, this->d}); - searchImplEx_( + searchImpl_( n, static_cast(vecs.data()), numeric_type, @@ -600,7 +600,7 @@ void GpuIndex::searchFromCpuPaged_( auto outIndicesSlice = outIndices.narrowOutermost(cur3, numToProcess); - searchImplEx_( + searchImpl_( numToProcess, static_cast(bufGpus[cur3BufIndex]->data()), numeric_type, diff --git a/faiss/gpu/GpuIndex.h b/faiss/gpu/GpuIndex.h index eb1f7b2d01..bdeb362d31 100644 --- a/faiss/gpu/GpuIndex.h +++ b/faiss/gpu/GpuIndex.h @@ -165,7 +165,7 @@ class GpuIndex : public faiss::Index { /// All data is guaranteed to be resident on our device virtual void addImpl_(idx_t n, const float* x, const idx_t* ids) = 0; - virtual void addImplEx_( + virtual void addImpl_( idx_t n, const void* x, NumericType numeric_type, @@ -187,7 +187,7 @@ class GpuIndex : public faiss::Index { idx_t* labels, const SearchParameters* params) const = 0; - virtual void searchImplEx_( + virtual void searchImpl_( idx_t n, const void* x, NumericType numeric_type, diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu index 6f129fe1d3..eada5936ff 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -152,7 +152,15 @@ 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::searchImplEx_( +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, NumericType numeric_type, @@ -240,7 +248,7 @@ void GpuIndexCagra::searchImpl_( float* distances, idx_t* labels, const SearchParameters* search_params) const { - searchImplEx_( + searchImpl_( n, static_cast(x), NumericType::Float32, diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index 1d60e3a72b..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_( @@ -294,7 +299,7 @@ struct GpuIndexCagra : public GpuIndex { float* distances, idx_t* labels, const SearchParameters* search_params) const override; - void searchImplEx_( + void searchImpl_( idx_t n, const void* x, NumericType numeric_type, diff --git a/faiss/gpu/GpuIndexFlat.cu b/faiss/gpu/GpuIndexFlat.cu index 78ed2410ad..a97bb07239 100644 --- a/faiss/gpu/GpuIndexFlat.cu +++ b/faiss/gpu/GpuIndexFlat.cu @@ -216,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, @@ -236,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 62dc049ded..4c541f9e36 100644 --- a/faiss/gpu/GpuIndexFlat.h +++ b/faiss/gpu/GpuIndexFlat.h @@ -123,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_( @@ -132,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_; From 00efb88e1d696f8123bbf0568874cf02abce019b Mon Sep 17 00:00:00 2001 From: jinsolp Date: Mon, 21 Jul 2025 23:36:25 +0000 Subject: [PATCH 16/16] call proper GpuIndex::func --- faiss/gpu/GpuIndexFlat.cu | 4 ++-- faiss/gpu/GpuIndexIVFFlat.cu | 2 +- faiss/gpu/GpuIndexIVFPQ.cu | 2 +- faiss/gpu/GpuIndexIVFScalarQuantizer.cu | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/faiss/gpu/GpuIndexFlat.cu b/faiss/gpu/GpuIndexFlat.cu index a97bb07239..c1b03456d0 100644 --- a/faiss/gpu/GpuIndexFlat.cu +++ b/faiss/gpu/GpuIndexFlat.cu @@ -168,7 +168,7 @@ void GpuIndexFlat::train(idx_t n, const float* x) { } void GpuIndexFlat::train(idx_t n, const void* x, NumericType numeric_type) { - Index::train(n, x, numeric_type); + GpuIndex::train(n, x, numeric_type); } void GpuIndexFlat::add(idx_t n, const float* x) { @@ -196,7 +196,7 @@ void GpuIndexFlat::add(idx_t n, const float* x) { } void GpuIndexFlat::add(idx_t n, const void* x, NumericType numeric_type) { - Index::add(n, x, numeric_type); + GpuIndex::add(n, x, numeric_type); } bool GpuIndexFlat::addImplRequiresIDs_() const { diff --git a/faiss/gpu/GpuIndexIVFFlat.cu b/faiss/gpu/GpuIndexIVFFlat.cu index 4a85795ac1..9d2d07ecda 100644 --- a/faiss/gpu/GpuIndexIVFFlat.cu +++ b/faiss/gpu/GpuIndexIVFFlat.cu @@ -325,7 +325,7 @@ void GpuIndexIVFFlat::train(idx_t n, const float* x) { } void GpuIndexIVFFlat::train(idx_t n, const void* x, NumericType numeric_type) { - Index::train(n, x, numeric_type); + GpuIndex::train(n, x, numeric_type); } void GpuIndexIVFFlat::setIndex_( diff --git a/faiss/gpu/GpuIndexIVFPQ.cu b/faiss/gpu/GpuIndexIVFPQ.cu index 5cd1ec1eb3..efa58c642f 100644 --- a/faiss/gpu/GpuIndexIVFPQ.cu +++ b/faiss/gpu/GpuIndexIVFPQ.cu @@ -490,7 +490,7 @@ void GpuIndexIVFPQ::train(idx_t n, const float* x) { } void GpuIndexIVFPQ::train(idx_t n, const void* x, NumericType numeric_type) { - Index::train(n, x, numeric_type); + GpuIndex::train(n, x, numeric_type); } void GpuIndexIVFPQ::setIndex_( diff --git a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu index d5b0e81e68..68670f1974 100644 --- a/faiss/gpu/GpuIndexIVFScalarQuantizer.cu +++ b/faiss/gpu/GpuIndexIVFScalarQuantizer.cu @@ -283,7 +283,7 @@ void GpuIndexIVFScalarQuantizer::train( idx_t n, const void* x, NumericType numeric_type) { - Index::train(n, x, numeric_type); + GpuIndex::train(n, x, numeric_type); } } // namespace gpu