facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add filtered search for cuVS indexes (#4858)

Browse Source 
Summary:
This PR adds support for pre-filtering search on cuVS indexes, using IDSelector. It also includes C++ tests for this feature.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4858

Reviewed By: limqiying

Differential Revision: D96034911

Pulled By: mnorris11

fbshipit-source-id: 1deaf78b9ba3d54b0f49925f6f0ef0b1f98e0f91
This commit is contained in:
Mickael Ide
2026-04-15 10:59:26 -07:00
committed by meta-codesync[bot]
parent 74dceee903
commit 3030fe0635
32 changed files with 681 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/gpu/GpuDistance.cu
+14 -2
View File
@@ -318,8 +318,14 @@ void bfKnn(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
cuvs::neighbors::brute_force::index<float> idx(
handle, index.view(), norms_view, distance, metric_arg);
cuvs::neighbors::brute_force::search_params search_params_bf;
cuvs::neighbors::brute_force::search(
handle, idx, search.view(), inds.view(), dists.view());
handle,
search_params_bf,
idx,
search.view(),
inds.view(),
dists.view());
} else {
auto index = raft::make_readonly_temporary_device_buffer<
const float,
@@ -357,8 +363,14 @@ void bfKnn(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
cuvs::neighbors::brute_force::index<float> idx(
handle, index.view(), norms_view, distance, metric_arg);
cuvs::neighbors::brute_force::search_params search_params_bf;
cuvs::neighbors::brute_force::search(
handle, idx, search.view(), inds.view(), dists.view());
handle,
search_params_bf,
idx,
search.view(),
inds.view(),
dists.view());
}
if (args.metric == MetricType::METRIC_Lp) {
faiss/gpu/GpuIndexBinaryCagra.cu
+2 -3
View File
@@ -119,8 +119,6 @@ void GpuIndexBinaryCagra::search(
return;
}
FAISS_THROW_IF_NOT_MSG(!params, "params not implemented");
// validateKSelect(k);
// The input vectors may be too large for the GPU, but we still
@@ -260,7 +258,8 @@ void GpuIndexBinaryCagra::searchImpl_(
params->hashmap_min_bitlen,
params->hashmap_max_fill_rate,
params->num_random_samplings,
params->seed);
params->seed,
params->sel);
if (not search_params) {
delete params;
faiss/gpu/GpuIndexCagra.cu
+6 -3
View File
@@ -228,7 +228,8 @@ void GpuIndexCagra::searchImpl_ex_(
params->hashmap_min_bitlen,
params->hashmap_max_fill_rate,
params->num_random_samplings,
params->seed);
params->seed,
params->sel);
} else if (numeric_type == NumericType::Float16) {
Tensor<half, 2, true> queries(
@@ -251,7 +252,8 @@ void GpuIndexCagra::searchImpl_ex_(
params->hashmap_min_bitlen,
params->hashmap_max_fill_rate,
params->num_random_samplings,
params->seed);
params->seed,
params->sel);
} else if (numeric_type == NumericType::Int8) {
Tensor<int8_t, 2, true> queries(
const_cast<int8_t*>(static_cast<const int8_t*>(x)),
@@ -274,7 +276,8 @@ void GpuIndexCagra::searchImpl_ex_(
params->hashmap_min_bitlen,
params->hashmap_max_fill_rate,
params->num_random_samplings,
params->seed);
params->seed,
params->sel);
} else {
FAISS_THROW_MSG("GpuIndexCagra::searchImpl_ unsupported data type");
}
faiss/gpu/GpuIndexFlat.cu
+9 -1
View File
@@ -224,8 +224,16 @@ void GpuIndexFlat::searchImpl_(
Tensor<float, 2, true> outDistances(distances, {n, k});
Tensor<idx_t, 2, true> outLabels(labels, {n, k});
const IDSelector* sel = params ? params->sel : nullptr;
data_->query(
queries, k, metric_type, metric_arg, outDistances, outLabels, true);
queries,
k,
metric_type,
metric_arg,
outDistances,
outLabels,
true,
sel);
}
void GpuIndexFlat::reconstruct(idx_t key, float* out) const {
faiss/gpu/GpuIndexIVF.cu
+3 -1
View File
@@ -342,8 +342,10 @@ void GpuIndexIVF::searchImpl_(
Tensor<float, 2, true> outDistances(distances, {n, k});
Tensor<idx_t, 2, true> outLabels(const_cast<idx_t*>(labels), {n, k});
const IDSelector* sel = params ? params->sel : nullptr;
baseIndex_->search(
quantizer, queries, use_nprobe, k, outDistances, outLabels);
quantizer, queries, use_nprobe, k, outDistances, outLabels, sel);
}
void GpuIndexIVF::search_preassigned(
faiss/gpu/impl/BinaryCuvsCagra.cu
+23 -2
View File
@@ -22,10 +22,12 @@
*/
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/utils/CuvsFilterConvert.h>
#include <faiss/gpu/utils/CuvsUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/gpu/impl/BinaryCuvsCagra.cuh>
#include <cuvs/core/bitset.hpp>
#include <cuvs/neighbors/cagra.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/core/device_resources.hpp>
@@ -202,7 +204,8 @@ void BinaryCuvsCagra::search(
idx_t hashmap_min_bitlen,
float hashmap_max_fill_rate,
idx_t num_random_samplings,
idx_t rand_xor_mask) {
idx_t rand_xor_mask,
const IDSelector* sel) {
const raft::device_resources& raft_handle =
resources_->getRaftHandleCurrentDevice();
idx_t numQueries = queries.getSize(0);
@@ -255,13 +258,31 @@ void BinaryCuvsCagra::search(
raft_handle, numQueries, k_);
auto distances_float_view = distances_float.view();
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_holder;
std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_holder =
cuvs::core::bitset<uint32_t, int64_t>(raft_handle, n_, false);
faiss::gpu::convert_to_bitset(resources_, *sel, bitset_holder->view());
bitset_filter.emplace(bitset_holder->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::cagra::search(
raft_handle,
search_pams,
*cuvs_index,
queries_view,
indices_copy.view(),
distances_float_view);
distances_float_view,
filter_ref);
thrust::copy(
raft::resource::get_thrust_policy(raft_handle),
faiss/gpu/impl/BinaryCuvsCagra.cuh
+3 -1
View File
@@ -31,6 +31,7 @@
#include <optional>
#include <faiss/MetricType.h>
#include <faiss/impl/IDSelector.h>
#include <cuvs/neighbors/cagra.hpp>
@@ -80,7 +81,8 @@ class BinaryCuvsCagra {
idx_t hashmap_min_bitlen,
float hashmap_max_fill_rate,
idx_t num_random_samplings,
idx_t rand_xor_mask);
idx_t rand_xor_mask,
const IDSelector* sel = nullptr);
void reset();
faiss/gpu/impl/CuvsCagra.cu
+24 -2
View File
@@ -22,10 +22,12 @@
*/
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/utils/CuvsFilterConvert.h>
#include <faiss/gpu/utils/CuvsUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/gpu/impl/CuvsCagra.cuh>
#include <cuvs/core/bitset.hpp>
#include <cuvs/neighbors/cagra.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/core/device_resources.hpp>
@@ -33,6 +35,7 @@
#include <thrust/copy.h>
#include <thrust/device_ptr.h>
#include <optional>
namespace faiss {
namespace gpu {
@@ -235,7 +238,8 @@ void CuvsCagra<data_t>::search(
idx_t hashmap_min_bitlen,
float hashmap_max_fill_rate,
idx_t num_random_samplings,
idx_t rand_xor_mask) {
idx_t rand_xor_mask,
const IDSelector* sel) {
const raft::device_resources& raft_handle =
resources_->getRaftHandleCurrentDevice();
idx_t numQueries = queries.getSize(0);
@@ -286,13 +290,31 @@ void CuvsCagra<data_t>::search(
auto indices_copy = raft::make_device_matrix<uint32_t, int64_t>(
raft_handle, numQueries, k_);
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_holder;
std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_holder =
cuvs::core::bitset<uint32_t, int64_t>(raft_handle, n_, false);
faiss::gpu::convert_to_bitset(resources_, *sel, bitset_holder->view());
bitset_filter.emplace(bitset_holder->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::cagra::search(
raft_handle,
search_pams,
*cuvs_index,
queries_view,
indices_copy.view(),
distances_view);
distances_view,
filter_ref);
thrust::copy(
raft::resource::get_thrust_policy(raft_handle),
indices_copy.data_handle(),
faiss/gpu/impl/CuvsCagra.cuh
+3 -1
View File
@@ -30,6 +30,7 @@
#include <optional>
#include <faiss/MetricType.h>
#include <faiss/impl/IDSelector.h>
#include <cuvs/neighbors/cagra.hpp>
#include <cuvs/neighbors/ivf_pq.hpp>
@@ -103,7 +104,8 @@ class CuvsCagra {
idx_t hashmap_min_bitlen,
float hashmap_max_fill_rate,
idx_t num_random_samplings,
idx_t rand_xor_mask);
idx_t rand_xor_mask,
const IDSelector* sel = nullptr);
void reset();
faiss/gpu/impl/CuvsFlatIndex.cu
+94 -21
View File
@@ -21,6 +21,7 @@
* limitations under the License.
*/
#include <faiss/gpu/utils/CuvsFilterConvert.h>
#include <faiss/gpu/utils/CuvsUtils.h>
#include <faiss/gpu/impl/CuvsFlatIndex.cuh>
#include <faiss/gpu/utils/ConversionOperators.cuh>
@@ -28,6 +29,8 @@
#include <optional>
#include <vector>
#include <cuvs/core/bitset.hpp>
#include <cuvs/distance/distance.h>
#include <cuvs/neighbors/brute_force.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/core/logger.hpp>
@@ -53,26 +56,22 @@ void CuvsFlatIndex::query(
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) {
/**
* cuVS doesn't yet support half-precision in bfknn.
* Use FlatIndex for float16 for now
*/
bool exactDistance,
const IDSelector* sel) {
if (useFloat16_) {
auto stream = resources_->getDefaultStreamCurrentDevice();
// We need to convert the input to float16 for comparison to ourselves
auto stream = resources_->getDefaultStreamCurrentDevice();
auto inputHalf = convertTensorTemporary<float, half, 2>(
resources_, stream, input);
FlatIndex::query(
CuvsFlatIndex::query(
inputHalf,
k,
metric,
metricArg,
outDistances,
outIndices,
exactDistance);
exactDistance,
sel);
} else {
raft::device_resources& handle =
resources_->getRaftHandleCurrentDevice();
@@ -99,7 +98,29 @@ void CuvsFlatIndex::query(
cuvs::neighbors::brute_force::index idx(
handle, index, norms_view, distance, metricArg);
cuvs::neighbors::brute_force::search(handle, idx, search, inds, dists);
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_cuvs;
std::optional<
cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter_cuvs;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_cuvs = cuvs::core::bitset<uint32_t, int64_t>(
handle, vectors_.getSize(0), false);
faiss::gpu::convert_to_bitset(
resources_, *sel, bitset_cuvs->view());
bitset_filter_cuvs.emplace(bitset_cuvs->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter_cuvs.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::brute_force::search_params search_params_bf;
cuvs::neighbors::brute_force::search(
handle, search_params_bf, idx, search, inds, dists, filter_ref);
if (metric == MetricType::METRIC_Lp) {
raft::linalg::unary_op(
@@ -126,18 +147,70 @@ void CuvsFlatIndex::query(
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) {
bool exactDistance,
const IDSelector* sel) {
FAISS_ASSERT(useFloat16_);
// FIXME: ref https://github.com/rapidsai/raft/issues/1280
FlatIndex::query(
vecs,
k,
metric,
metricArg,
outDistances,
outIndices,
exactDistance);
raft::device_resources& handle = resources_->getRaftHandleCurrentDevice();
auto index = raft::make_device_matrix_view<const half, int64_t>(
vectorsHalf_.data(),
vectorsHalf_.getSize(0),
vectorsHalf_.getSize(1));
auto search = raft::make_device_matrix_view<const half, int64_t>(
vecs.data(), vecs.getSize(0), vecs.getSize(1));
auto inds = raft::make_device_matrix_view<idx_t, int64_t>(
outIndices.data(), outIndices.getSize(0), outIndices.getSize(1));
auto dists = raft::make_device_matrix_view<float, int64_t>(
outDistances.data(),
outDistances.getSize(0),
outDistances.getSize(1));
auto distance = metricFaissToCuvs(metric, exactDistance);
std::optional<raft::device_vector_view<const float, int64_t>> norms_view =
raft::make_device_vector_view(norms_.data(), norms_.getSize(0));
cuvs::neighbors::brute_force::index<half, float> idx(
handle, index, norms_view, distance, metricArg);
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_cuvs;
std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter_cuvs;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_cuvs = cuvs::core::bitset<uint32_t, int64_t>(
handle, vectorsHalf_.getSize(0), false);
faiss::gpu::convert_to_bitset(resources_, *sel, bitset_cuvs->view());
bitset_filter_cuvs.emplace(bitset_cuvs->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter_cuvs.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::brute_force::search_params search_params_bf;
cuvs::neighbors::brute_force::search(
handle, search_params_bf, idx, search, inds, dists, filter_ref);
if (metric == MetricType::METRIC_Lp) {
raft::linalg::unary_op(
handle,
raft::make_const_mdspan(dists),
dists,
[metricArg] __device__(const float& a) {
return powf(a, metricArg);
});
} else if (metric == MetricType::METRIC_JensenShannon) {
raft::linalg::unary_op(
handle,
raft::make_const_mdspan(dists),
dists,
[] __device__(const float& a) { return powf(a, 2); });
}
}
} // namespace gpu
faiss/gpu/impl/CuvsFlatIndex.cuh
+5 -2
View File
@@ -25,6 +25,7 @@
#include <faiss/MetricType.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/gpu/impl/FlatIndex.cuh>
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/DeviceVector.cuh>
@@ -55,7 +56,8 @@ class CuvsFlatIndex : public FlatIndex {
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) override;
bool exactDistance,
const IDSelector* sel) override;
void query(
Tensor<half, 2, true>& vecs,
@@ -64,7 +66,8 @@ class CuvsFlatIndex : public FlatIndex {
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) override;
bool exactDistance,
const IDSelector* sel) override;
};
} // namespace gpu
faiss/gpu/impl/CuvsIVFFlat.cu
+24 -2
View File
@@ -26,14 +26,17 @@
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/utils/CuvsFilterConvert.h>
#include <faiss/gpu/utils/CuvsUtils.h>
#include <faiss/gpu/impl/CuvsIVFFlat.cuh>
#include <faiss/gpu/impl/FlatIndex.cuh>
#include <faiss/gpu/impl/IVFFlat.cuh>
#include <faiss/gpu/utils/Transpose.cuh>
#include <cuvs/core/bitset.hpp>
#include <cuvs/neighbors/common.hpp>
#include <cuvs/neighbors/ivf_flat.hpp>
#include <optional>
#include <raft/linalg/map.cuh>
#include <raft/linalg/norm.cuh>
@@ -100,7 +103,8 @@ void CuvsIVFFlat::search(
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) {
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel) {
/// NB: The coarse quantizer is ignored here. The user is assumed to have
/// called updateQuantizer() to modify the cuVS index if the quantizer was
/// modified externally
@@ -127,13 +131,31 @@ void CuvsIVFFlat::search(
auto out_dists_view = raft::make_device_matrix_view<float, idx_t>(
outDistances.data(), (idx_t)numQueries, (idx_t)k_);
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_cuvs;
std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter_cuvs;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_cuvs = cuvs::core::bitset<uint32_t, int64_t>(
raft_handle, cuvs_index->size(), false);
faiss::gpu::convert_to_bitset(resources_, *sel, bitset_cuvs->view());
bitset_filter_cuvs.emplace(bitset_cuvs->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter_cuvs.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::ivf_flat::search(
raft_handle,
pams,
*cuvs_index,
queries_view,
out_inds_view,
out_dists_view);
out_dists_view,
filter_ref);
/// Identify NaN rows and mask their nearest neighbors
auto nan_flag = raft::make_device_vector<bool>(raft_handle, numQueries);
faiss/gpu/impl/CuvsIVFFlat.cuh
+3 -1
View File
@@ -24,6 +24,7 @@
#pragma once
#include <faiss/impl/CodePacker.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/gpu/impl/GpuScalarQuantizer.cuh>
#include <faiss/gpu/impl/IVFFlat.cuh>
@@ -63,7 +64,8 @@ class CuvsIVFFlat : public IVFFlat {
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) override;
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel = nullptr) override;
/// Performs search when we are already given the IVF cells to look at
/// (GpuIndexIVF::search_preassigned implementation)
faiss/gpu/impl/CuvsIVFPQ.cu
+24 -2
View File
@@ -23,17 +23,20 @@
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/utils/CuvsFilterConvert.h>
#include <faiss/gpu/utils/CuvsUtils.h>
#include <faiss/gpu/impl/CuvsIVFPQ.cuh>
#include <faiss/gpu/impl/FlatIndex.cuh>
#include <faiss/gpu/utils/Transpose.cuh>
#include <cuvs/core/bitset.hpp>
#include <cuvs/neighbors/common.hpp>
#include <cuvs/neighbors/ivf_pq.hpp>
#include <raft/linalg/map.cuh>
#include <limits>
#include <memory>
#include <optional>
namespace faiss {
namespace gpu {
@@ -359,7 +362,8 @@ void CuvsIVFPQ::search(
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) {
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel) {
uint32_t numQueries = queries.getSize(0);
uint32_t cols = queries.getSize(1);
idx_t k_ = std::min(static_cast<idx_t>(k), cuvs_index->size());
@@ -383,13 +387,31 @@ void CuvsIVFPQ::search(
auto out_dists_view = raft::make_device_matrix_view<float, idx_t>(
outDistances.data(), (idx_t)numQueries, (idx_t)k_);
std::optional<cuvs::core::bitset<uint32_t, int64_t>> bitset_cuvs;
std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
bitset_filter_cuvs;
cuvs::neighbors::filtering::none_sample_filter none_filter;
if (sel) {
bitset_cuvs = cuvs::core::bitset<uint32_t, int64_t>(
raft_handle, cuvs_index->size(), false);
faiss::gpu::convert_to_bitset(resources_, *sel, bitset_cuvs->view());
bitset_filter_cuvs.emplace(bitset_cuvs->view());
}
const cuvs::neighbors::filtering::base_filter& filter_ref = sel
? static_cast<const cuvs::neighbors::filtering::base_filter&>(
bitset_filter_cuvs.value())
: static_cast<const cuvs::neighbors::filtering::base_filter&>(
none_filter);
cuvs::neighbors::ivf_pq::search(
raft_handle,
pams,
*cuvs_index,
queries_view,
out_inds_view,
out_dists_view);
out_dists_view,
filter_ref);
/// Identify NaN rows and mask their nearest neighbors
auto nan_flag = raft::make_device_vector<bool>(raft_handle, numQueries);
faiss/gpu/impl/CuvsIVFPQ.cuh
+2 -1
View File
@@ -76,7 +76,8 @@ class CuvsIVFPQ : public IVFPQ {
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) override;
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel = nullptr) override;
/// Performs search when we are already given the IVF cells to look at
/// (GpuIndexIVF::search_preassigned implementation)
faiss/gpu/impl/FlatIndex.cu
+6 -3
View File
@@ -116,7 +116,8 @@ void FlatIndex::query(
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) {
bool exactDistance,
const IDSelector* sel) {
auto stream = resources_->getDefaultStreamCurrentDevice();
if (useFloat16_) {
@@ -130,7 +131,8 @@ void FlatIndex::query(
metricArg,
outDistances,
outIndices,
exactDistance);
exactDistance,
sel);
} else {
bfKnnOnDevice(
resources_,
@@ -157,7 +159,8 @@ void FlatIndex::query(
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance) {
bool exactDistance,
const IDSelector* sel) {
FAISS_ASSERT(useFloat16_);
bfKnnOnDevice(
faiss/gpu/impl/FlatIndex.cuh
+5 -2
View File
@@ -25,6 +25,7 @@
#include <faiss/MetricType.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/DeviceVector.cuh>
@@ -67,7 +68,8 @@ class FlatIndex {
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance);
bool exactDistance,
const IDSelector* sel = nullptr);
virtual void query(
Tensor<half, 2, true>& vecs,
@@ -76,7 +78,8 @@ class FlatIndex {
float metricArg,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices,
bool exactDistance);
bool exactDistance,
const IDSelector* sel = nullptr);
/// Compute residual for set of vectors
void computeResidual(
faiss/gpu/impl/IVFBase.cuh
+3 -1
View File
@@ -10,6 +10,7 @@
#include <faiss/Index.h>
#include <faiss/MetricType.h>
#include <faiss/gpu/GpuIndicesOptions.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/DeviceVector.cuh>
#include <memory>
@@ -95,7 +96,8 @@ class IVFBase {
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) = 0;
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel = nullptr) = 0;
/// Performs search when we are already given the IVF cells to look at
/// (GpuIndexIVF::search_preassigned implementation)
faiss/gpu/impl/IVFFlat.cu
+2 -1
View File
@@ -169,7 +169,8 @@ void IVFFlat::search(
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) {
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel) {
auto stream = resources_->getDefaultStreamCurrentDevice();
// These are caught at a higher level
faiss/gpu/impl/IVFFlat.cuh
+2 -1
View File
@@ -37,7 +37,8 @@ class IVFFlat : public IVFBase {
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) override;
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel = nullptr) override;
/// Performs search when we are already given the IVF cells to look at
/// (GpuIndexIVF::search_preassigned implementation)
faiss/gpu/impl/IVFPQ.cu
+2 -1
View File
@@ -494,7 +494,8 @@ void IVFPQ::search(
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) {
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel) {
// These are caught at a higher level
FAISS_ASSERT(nprobe <= GPU_MAX_SELECTION_K);
FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
faiss/gpu/impl/IVFPQ.cuh
+2 -1
View File
@@ -53,7 +53,8 @@ class IVFPQ : public IVFBase {
int nprobe,
int k,
Tensor<float, 2, true>& outDistances,
Tensor<idx_t, 2, true>& outIndices) override;
Tensor<idx_t, 2, true>& outIndices,
const IDSelector* sel = nullptr) override;
/// Performs search when we are already given the IVF cells to look at
/// (GpuIndexIVF::search_preassigned implementation)
faiss/gpu/test/TestGpuFilterConvert.cu
+25 -18
View File
@@ -53,6 +53,8 @@ struct Options {
};
void run_complex() {
using bitset_t = uint32_t;
using indexing_t = int64_t;
Options spec;
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
@@ -77,7 +79,7 @@ void run_complex() {
auto or_selector = faiss::IDSelectorOr(&range_selector, &array_selector);
auto bitmap_faiss_cpu = std::vector<uint8_t>((spec.bitset_len + 8) / 8);
for (uint32_t i = 0; i < bitmap_faiss_cpu.size(); i++) {
for (indexing_t i = 0; i < bitmap_faiss_cpu.size(); i++) {
bitmap_faiss_cpu[i] = (uint8_t)faiss::gpu::randVal(0, 255);
}
auto bitmap_selector =
@@ -88,19 +90,19 @@ void run_complex() {
faiss::IDSelectorXOr(&or_selector, &not_bitmap_selector);
// convert to cuVS bitset
auto bitset = cuvs::core::bitset<uint32_t, uint32_t>(
auto bitset = cuvs::core::bitset<bitset_t, indexing_t>(
raft_handle, spec.bitset_len, false);
faiss::gpu::convert_to_bitset(gpuRes.get(), xor_selector, bitset.view());
// verify
auto bitset_converted_cpu =
raft::make_host_vector<uint32_t, uint32_t>(bitset.n_elements());
raft::make_host_vector<bitset_t, indexing_t>(bitset.n_elements());
auto bitset_converted_cpu_view =
cuvs::core::bitset_view<uint32_t, uint32_t>(
cuvs::core::bitset_view<bitset_t, indexing_t>(
bitset_converted_cpu.data_handle(), spec.bitset_len);
raft::copy(raft_handle, bitset_converted_cpu.view(), bitset.to_mdspan());
raft::resource::sync_stream(raft_handle);
for (uint32_t i = 0; i < spec.bitset_len; i++) {
for (indexing_t i = 0; i < spec.bitset_len; i++) {
if (bitset_converted_cpu_view.test(i) != xor_selector.is_member(i)) {
ASSERT_TRUE(
testing::AssertionFailure()
@@ -120,23 +122,24 @@ void run_range() {
gpuRes->getRaftHandleCurrentDevice();
// take random imin and imax, check all ids
using bitset_t = uint32_t;
using indexing_t = int64_t;
auto imin = faiss::gpu::randVal(0, spec.bitset_len - 2);
auto imax = faiss::gpu::randVal(1, spec.bitset_len - 1);
if (imin > imax)
std::swap(imin, imax);
auto selector = faiss::IDSelectorRange(imin, imax);
auto bitset = cuvs::core::bitset<bitset_t, uint32_t>(
auto bitset = cuvs::core::bitset<bitset_t, indexing_t>(
raft_handle, spec.bitset_len, false);
auto nbits = sizeof(bitset_t) * 8;
faiss::gpu::convert_to_bitset(gpuRes.get(), selector, bitset.view());
auto bitset_converted_cpu =
raft::make_host_vector<bitset_t, uint32_t>(bitset.n_elements());
raft::make_host_vector<bitset_t, indexing_t>(bitset.n_elements());
raft::copy(raft_handle, bitset_converted_cpu.view(), bitset.to_mdspan());
raft::resource::sync_stream(raft_handle);
auto bitset_view_cpu = cuvs::core::bitset_view<bitset_t, uint32_t>(
auto bitset_view_cpu = cuvs::core::bitset_view<bitset_t, indexing_t>(
bitset_converted_cpu.data_handle(), spec.bitset_len);
for (uint64_t i = 0; i < spec.bitset_len; i++) {
for (indexing_t i = 0; i < spec.bitset_len; i++) {
if (bitset_view_cpu.test(i) != selector.is_member(i)) {
ASSERT_TRUE(
testing::AssertionFailure()
@@ -154,28 +157,30 @@ void run_bitmap() {
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
using bitset_t = uint32_t;
using indexing_t = int64_t;
auto gpuRes = res.getResources();
const raft::device_resources& raft_handle =
gpuRes->getRaftHandleCurrentDevice();
// generate random bitmap selector
auto bitmap_faiss_cpu = std::vector<uint8_t>((spec.bitset_len + 8) / 8);
for (uint32_t i = 0; i < bitmap_faiss_cpu.size(); i++) {
for (indexing_t i = 0; i < bitmap_faiss_cpu.size(); i++) {
bitmap_faiss_cpu[i] = (uint8_t)faiss::gpu::randVal(0, 255);
}
auto bitmap_selector =
faiss::IDSelectorBitmap(spec.bitset_len, bitmap_faiss_cpu.data());
auto bitset = cuvs::core::bitset<uint32_t, uint32_t>(
auto bitset = cuvs::core::bitset<bitset_t, indexing_t>(
raft_handle, spec.bitset_len, false);
faiss::gpu::convert_to_bitset(gpuRes.get(), bitmap_selector, bitset.view());
auto bitset_converted_cpu =
raft::make_host_vector<uint32_t, uint32_t>(bitset.n_elements());
raft::make_host_vector<bitset_t, indexing_t>(bitset.n_elements());
raft::copy(raft_handle, bitset_converted_cpu.view(), bitset.to_mdspan());
raft::resource::sync_stream(raft_handle);
auto bitset_converted_cpu_view =
cuvs::core::bitset_view<uint32_t, uint32_t>(
cuvs::core::bitset_view<bitset_t, indexing_t>(
bitset_converted_cpu.data_handle(), spec.bitset_len);
for (uint32_t i = 0; i < spec.bitset_len; i++) {
for (indexing_t i = 0; i < spec.bitset_len; i++) {
if (bitset_converted_cpu_view.test(i) != bitmap_selector.is_member(i)) {
ASSERT_TRUE(
testing::AssertionFailure()
@@ -189,6 +194,8 @@ void run_bitmap() {
}
void run_array() {
using bitset_t = uint32_t;
using indexing_t = int64_t;
Options spec;
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
@@ -203,18 +210,18 @@ void run_array() {
}
auto array_selector =
faiss::IDSelectorArray(n, array_selector_indices.data());
auto bitset = cuvs::core::bitset<uint32_t, uint32_t>(
auto bitset = cuvs::core::bitset<bitset_t, indexing_t>(
raft_handle, spec.bitset_len, false);
faiss::gpu::convert_to_bitset(gpuRes.get(), array_selector, bitset.view());
auto bitset_converted_cpu =
raft::make_host_vector<uint32_t, uint32_t>(bitset.n_elements());
raft::make_host_vector<bitset_t, indexing_t>(bitset.n_elements());
raft::copy(raft_handle, bitset_converted_cpu.view(), bitset.to_mdspan());
raft::resource::sync_stream(raft_handle);
auto bitset_converted_cpu_view =
cuvs::core::bitset_view<uint32_t, uint32_t>(
cuvs::core::bitset_view<bitset_t, indexing_t>(
bitset_converted_cpu.data_handle(), spec.bitset_len);
for (uint32_t i = 0; i < spec.bitset_len; i++) {
for (indexing_t i = 0; i < spec.bitset_len; i++) {
if (bitset_converted_cpu_view.test(i) != array_selector.is_member(i)) {
ASSERT_TRUE(
testing::AssertionFailure()
faiss/gpu/test/TestGpuIndexBinaryCagra.cu
+69
View File
@@ -26,6 +26,7 @@
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/distances.h>
#include <cstddef>
#include <faiss/gpu/utils/CopyUtils.cuh>
@@ -432,6 +433,74 @@ TEST(TestGpuIndexBinaryCagra, CopyFrom_ITERATIVE_SEARCH) {
copyFromTest(faiss::gpu::graph_build_algo::ITERATIVE_SEARCH, 0.98);
}
void testIDSelectorBinaryCagra(std::string selectorName) {
Options opt;
auto trainVecs = faiss::gpu::randBinaryVecs(opt.numTrain, opt.dim);
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
faiss::gpu::GpuIndexCagraConfig config;
config.device = opt.device;
config.graph_degree = opt.graphDegree;
config.intermediate_graph_degree = opt.intermediateGraphDegree;
config.build_algo = faiss::gpu::graph_build_algo::NN_DESCENT;
config.nn_descent_niter = 20;
faiss::gpu::GpuIndexBinaryCagra gpuIndex(&res, opt.dim, config);
gpuIndex.train(opt.numTrain, trainVecs.data());
auto queryVecs = faiss::gpu::randBinaryVecs(opt.numQuery, opt.dim);
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numTrain);
faiss::gpu::SearchParametersCagra search_params;
// TODO: For CI test purposes only, remove this
// for (auto& [selectorName, selector] : selector_struct.selector_map) {
auto selector = selector_struct.selector_map[selectorName].get();
search_params.sel = selector;
std::vector<int> distances(opt.numQuery * opt.k, 0);
std::vector<faiss::idx_t> labels(opt.numQuery * opt.k, -1);
gpuIndex.search(
opt.numQuery,
queryVecs.data(),
opt.k,
distances.data(),
labels.data(),
&search_params);
for (int i = 0; i < opt.numQuery * opt.k; ++i) {
if (labels[i] >= 0) {
EXPECT_TRUE(selector->is_member(labels[i]))
<< "Label " << labels[i] << " @ " << i << " not in "
<< selectorName << " selector";
}
}
//}
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Range) {
testIDSelectorBinaryCagra("Range");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Array) {
testIDSelectorBinaryCagra("Array");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Batch) {
testIDSelectorBinaryCagra("Batch");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Bitmap) {
testIDSelectorBinaryCagra("Bitmap");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Not) {
testIDSelectorBinaryCagra("Not");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_And) {
testIDSelectorBinaryCagra("And");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_Or) {
testIDSelectorBinaryCagra("Or");
}
TEST(TestGpuIndexBinaryCagra, IDSelector_XOr) {
testIDSelectorBinaryCagra("XOr");
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
faiss/gpu/test/TestGpuIndexCagra.cu
+48
View File
@@ -28,6 +28,7 @@
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/distances.h>
#include <cstddef>
#include <faiss/gpu/utils/CopyUtils.cuh>
@@ -895,6 +896,53 @@ TEST(TestGpuIndexCagra, Float16_CopyFrom_IP) {
copyFromTestFP16(faiss::METRIC_INNER_PRODUCT, 0.98);
}
void testIDSelectorCagra(faiss::MetricType metricType) {
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
if (metricType == faiss::METRIC_INNER_PRODUCT) {
faiss::fvec_renorm_L2(opt.numTrain, opt.dim, trainVecs.data());
}
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
faiss::gpu::GpuIndexCagraConfig config;
config.device = opt.device;
config.graph_degree = opt.graphDegree;
config.intermediate_graph_degree = opt.intermediateGraphDegree;
// Use only IVF_PQ to avoid NN_DESCENT + IP combination
config.build_algo = faiss::gpu::graph_build_algo::IVF_PQ;
faiss::gpu::GpuIndexCagra gpuIndex(&res, opt.dim, metricType, config);
gpuIndex.train(opt.numTrain, trainVecs.data());
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
if (metricType == faiss::METRIC_INNER_PRODUCT) {
faiss::fvec_renorm_L2(opt.numQuery, opt.dim, queryVecs.data());
}
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numTrain);
faiss::gpu::SearchParametersCagra search_params;
for (auto& [selectorName, selector] : selector_struct.selector_map) {
search_params.sel = selector.get();
faiss::gpu::testIDSelectorSearch(
&gpuIndex,
&search_params,
queryVecs,
opt.numQuery,
opt.k,
selectorName);
}
}
TEST(TestGpuIndexCagra, IDSelector_L2) {
testIDSelectorCagra(faiss::METRIC_L2);
}
TEST(TestGpuIndexCagra, IDSelector_IP) {
testIDSelectorCagra(faiss::METRIC_INNER_PRODUCT);
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
faiss/gpu/test/TestGpuIndexFlat.cpp
+46
View File
@@ -11,6 +11,7 @@
#include <faiss/gpu/impl/IndexUtils.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/impl/IDSelector.h>
#include <gtest/gtest.h>
#include <sstream>
#include <unordered_map>
@@ -761,6 +762,51 @@ TEST(TestCuvsGpuIndexFlat, SearchAndReconstruct) {
}
#endif
void testIDSelectorFlat(faiss::MetricType metricType) {
int numAdd = faiss::gpu::randVal(2000, 5000);
int dim = faiss::gpu::randVal(64, 200);
int numQuery = faiss::gpu::randVal(32, 100);
int k = faiss::gpu::randVal(10, 30);
int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
auto queryVecs = faiss::gpu::randVecs(numQuery, dim);
faiss::gpu::GpuIndexFlatConfig config;
config.device = device;
config.use_cuvs = true;
for (bool useFloat16 : {false, true}) {
config.useFloat16 = useFloat16;
faiss::gpu::GpuIndexFlat gpuIndex(&res, dim, metricType, config);
std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
gpuIndex.add(numAdd, addVecs.data());
faiss::gpu::TestIDSelectorStruct selector_struct(numAdd);
faiss::SearchParameters search_params;
for (auto& [selectorName, selector] : selector_struct.selector_map) {
search_params.sel = selector.get();
faiss::gpu::testIDSelectorSearch(
&gpuIndex,
&search_params,
queryVecs,
numQuery,
k,
selectorName);
}
}
}
#if defined USE_NVIDIA_CUVS
TEST(TestCuvsGpuIndexFlat, IDSelector_L2) {
testIDSelectorFlat(faiss::METRIC_L2);
}
TEST(TestCuvsGpuIndexFlat, IDSelector_IP) {
testIDSelectorFlat(faiss::METRIC_INNER_PRODUCT);
}
#endif
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
faiss/gpu/test/TestGpuIndexIVFFlat.cpp
+48
View File
@@ -28,6 +28,7 @@
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/distances.h>
#include <gtest/gtest.h>
#include <cmath>
@@ -917,6 +918,53 @@ TEST(TestGpuIndexIVFFlat, Reconstruct_n) {
EXPECT_EQ(gpuVals, cpuVals);
}
void testIDSelectorIVFFlat(faiss::MetricType metricType) {
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
faiss::gpu::GpuIndexIVFFlatConfig config;
config.device = opt.device;
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
config.use_cuvs = true;
faiss::gpu::GpuIndexIVFFlat gpuIndex(
&res, opt.dim, opt.numCentroids, metricType, config);
gpuIndex.nprobe = opt.nprobe;
gpuIndex.train(opt.numTrain, trainVecs.data());
gpuIndex.add(opt.numAdd, addVecs.data());
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
faiss::SearchParametersIVF search_params;
search_params.nprobe = opt.nprobe;
for (auto& [selectorName, selector] : selector_struct.selector_map) {
search_params.sel = selector.get();
faiss::gpu::testIDSelectorSearch(
&gpuIndex,
&search_params,
queryVecs,
opt.numQuery,
opt.k,
selectorName);
}
}
#if defined USE_NVIDIA_CUVS
TEST(TestCuvsGpuIndexIVFFlat, IDSelector_L2) {
testIDSelectorIVFFlat(faiss::METRIC_L2);
}
TEST(TestCuvsGpuIndexIVFFlat, IDSelector_IP) {
testIDSelectorIVFFlat(faiss::METRIC_INNER_PRODUCT);
}
#endif
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
faiss/gpu/test/TestGpuIndexIVFPQ.cpp
+58
View File
@@ -6,11 +6,13 @@
*/
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/gpu/GpuIndexIVFPQ.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/gpu/utils/DeviceUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/distances.h>
#include <gtest/gtest.h>
#include <cmath>
@@ -896,6 +898,62 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
#endif
}
void testIDSelectorIVFPQ(faiss::MetricType metricType) {
Options opt;
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
faiss::IndexFlatL2 quantizerL2(opt.dim);
faiss::IndexFlatIP quantizerIP(opt.dim);
faiss::Index* quantizer = metricType == faiss::METRIC_L2
? (faiss::Index*)&quantizerL2
: (faiss::Index*)&quantizerIP;
faiss::IndexIVFPQ cpuIndex(
quantizer, opt.dim, opt.numCentroids, opt.codes, opt.bitsPerCode);
cpuIndex.metric_type = metricType;
cpuIndex.train(opt.numTrain, trainVecs.data());
cpuIndex.add(opt.numAdd, addVecs.data());
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
faiss::gpu::GpuIndexIVFPQConfig config;
config.device = opt.device;
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
config.interleavedLayout = true;
config.use_cuvs = true;
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
gpuIndex.nprobe = opt.nprobe;
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
faiss::SearchParametersIVF search_params;
search_params.nprobe = opt.nprobe;
for (auto& [selectorName, selector] : selector_struct.selector_map) {
search_params.sel = selector.get();
faiss::gpu::testIDSelectorSearch(
&gpuIndex,
&search_params,
queryVecs,
opt.numQuery,
opt.k,
selectorName);
}
}
#if defined USE_NVIDIA_CUVS
TEST(TestCuvsGpuIndexIVFPQ, IDSelector_L2) {
testIDSelectorIVFPQ(faiss::METRIC_L2);
}
TEST(TestCuvsGpuIndexIVFPQ, IDSelector_IP) {
testIDSelectorIVFPQ(faiss::METRIC_INNER_PRODUCT);
}
#endif
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
faiss/gpu/test/TestUtils.cpp
+72
View File
@@ -6,6 +6,7 @@
*/
#include <faiss/gpu/test/TestUtils.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/random.h>
#include <gtest/gtest.h>
#include <time.h>
@@ -429,5 +430,76 @@ void compareLists(
}
}
TestIDSelectorStruct::TestIDSelectorStruct(int numAdd) {
// Range selector [20%, 80%] of the database
size_t min_id = numAdd / 5;
size_t max_id = numAdd * 4 / 5;
selector_map["Range"] =
std::make_unique<faiss::IDSelectorRange>(min_id, max_id);
// Array selector (every 3rd element)
array_ids.clear();
for (int i = 0; i < numAdd; i += 3) {
array_ids.push_back(i);
}
selector_map["Array"] = std::make_unique<faiss::IDSelectorArray>(
array_ids.size(), array_ids.data());
// Batch selector (every 5th element)
batch_ids.clear();
for (int i = 1; i < numAdd; i += 5) {
batch_ids.push_back(i);
}
selector_map["Batch"] = std::make_unique<faiss::IDSelectorBatch>(
batch_ids.size(), batch_ids.data());
// Bitmap selector (every 4th element selected)
size_t bitmap_size = (numAdd + 7) / 8;
bitmap.resize(bitmap_size, 0);
for (int i = 0; i < numAdd; i += 4) {
int byte_idx = i / 8;
int bit_idx = i % 8;
bitmap[byte_idx] |= (1 << bit_idx);
}
selector_map["Bitmap"] =
std::make_unique<faiss::IDSelectorBitmap>(numAdd, bitmap.data());
selector_map["Not"] =
std::make_unique<faiss::IDSelectorNot>(selector_map["Range"].get());
selector_map["And"] = std::make_unique<faiss::IDSelectorAnd>(
selector_map["Range"].get(), selector_map["Array"].get());
selector_map["Or"] = std::make_unique<faiss::IDSelectorOr>(
selector_map["Range"].get(), selector_map["Batch"].get());
selector_map["XOr"] = std::make_unique<faiss::IDSelectorXOr>(
selector_map["Not"].get(), selector_map["Array"].get());
}
void testIDSelectorSearch(
faiss::Index* index,
faiss::SearchParameters* search_params,
const std::vector<float>& queryVecs,
int numQuery,
int k,
const std::string& selectorName) {
FAISS_ASSERT(search_params && search_params->sel);
std::vector<float> distances(numQuery * k, 0);
std::vector<faiss::idx_t> labels(numQuery * k, -1);
index->search(
numQuery,
queryVecs.data(),
k,
distances.data(),
labels.data(),
search_params);
faiss::IDSelector* selector = search_params->sel;
for (int i = 0; i < numQuery * k; ++i) {
if (labels[i] >= 0) {
EXPECT_TRUE(selector->is_member(labels[i]))
<< "Label " << labels[i] << " @ " << i << " not in "
<< selectorName << " selector";
}
}
}
} // namespace gpu
} // namespace faiss
faiss/gpu/test/TestUtils.h
+23
View File
@@ -13,6 +13,7 @@
#include <gtest/gtest.h>
#include <cstring>
#include <initializer_list>
#include <map>
#include <memory>
#include <string>
#include <vector>
@@ -141,5 +142,27 @@ void testIVFEquality(A& cpuIndex, B& gpuIndex) {
}
}
/// Run search with the given search_params and verify all returned labels are
/// members of the selector. Works with any Index that supports search with
/// SearchParameters.
void testIDSelectorSearch(
faiss::Index* index,
faiss::SearchParameters* search_params,
const std::vector<float>& queryVecs,
int numQuery,
int k,
const std::string& selectorName);
// Structure to hold all IDSelector instances
struct TestIDSelectorStruct {
// Storage for selectors that need it
std::vector<faiss::idx_t> array_ids;
std::vector<faiss::idx_t> batch_ids;
std::vector<uint8_t> bitmap;
std::map<std::string, std::unique_ptr<faiss::IDSelector>> selector_map;
explicit TestIDSelectorStruct(int numAdd);
};
} // namespace gpu
} // namespace faiss
faiss/gpu/utils/CuvsFilterConvert.cu
+30 -29
View File
@@ -36,13 +36,13 @@ namespace faiss::gpu {
template <typename bitset_t>
RAFT_KERNEL set_range_kernel(
bitset_t* bitset_data,
uint32_t imin,
uint32_t imax,
uint32_t n_elements_to_set) {
uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
int64_t imin,
int64_t imax,
int64_t n_elements_to_set) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t nbits = sizeof(bitset_t) * 8;
uint32_t current_index = (imin / nbits) + idx;
int64_t current_index = (imin / nbits) + idx;
bitset_t mask = 0;
if (idx < n_elements_to_set) {
if (n_elements_to_set == 1) {
@@ -70,7 +70,7 @@ RAFT_KERNEL set_range_kernel(
void convert_to_bitset_range(
raft::resources const& res,
const faiss::IDSelectorRange& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset) {
cuvs::core::bitset_view<uint32_t, int64_t> bitset) {
RAFT_EXPECTS(
bitset.size() >= selector.imax,
"IDSelectorRange is out of range for the given bitset");
@@ -79,10 +79,10 @@ void convert_to_bitset_range(
if (original_nbits == 0) {
original_nbits = nbits;
}
uint32_t imin = selector.imin;
uint32_t imax = selector.imax;
int64_t imin = selector.imin;
int64_t imax = selector.imax;
uint32_t n_elements_to_set = 1 + (imax + original_nbits) / original_nbits;
int64_t n_elements_to_set = 1 + (imax + original_nbits) / original_nbits;
n_elements_to_set -= (imin + original_nbits) / original_nbits;
auto stream = raft::resource::get_cuda_stream(res);
@@ -107,14 +107,14 @@ void convert_to_bitset_range(
void convert_to_bitset_array(
raft::resources const& res,
const faiss::IDSelectorArray& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset) {
uint32_t n = selector.n;
cuvs::core::bitset_view<uint32_t, int64_t> bitset) {
int64_t n = selector.n;
auto d_indexes_to_set =
raft::make_device_vector<faiss::idx_t, uint32_t>(res, n);
raft::make_device_vector<faiss::idx_t, int64_t>(res, n);
raft::copy(
res,
d_indexes_to_set.view(),
raft::make_host_vector_view<const faiss::idx_t, uint32_t>(
raft::make_host_vector_view<const faiss::idx_t, int64_t>(
selector.ids, n));
thrust::for_each_n(
raft::resource::get_thrust_policy(res),
@@ -128,13 +128,13 @@ void convert_to_bitset_array(
RAFT_KERNEL set_bitmap_kernel(
uint32_t* new_bitset_data,
uint8_t* original_bitmap_data,
uint32_t n_elements,
uint32_t bitset_original_nbits) {
uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t bit_index = 0;
uint32_t bit_offset = 0;
int64_t n_elements,
int64_t bitset_original_nbits) {
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
int64_t bit_index = 0;
int64_t bit_offset = 0;
raft::core::compute_original_nbits_position(
uint32_t{8}, bitset_original_nbits, idx * 8, bit_index, bit_offset);
int64_t{8}, bitset_original_nbits, idx * 8, bit_index, bit_offset);
if (idx < n_elements) {
uint32_t mask = original_bitmap_data[idx];
atomicOr(&new_bitset_data[bit_index], mask << bit_offset);
@@ -144,24 +144,25 @@ RAFT_KERNEL set_bitmap_kernel(
void convert_to_bitset_bitmap(
raft::resources const& res,
const faiss::IDSelectorBitmap& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset) {
uint32_t n = selector.n;
cuvs::core::bitset_view<uint32_t, int64_t> bitset) {
auto n = selector.n;
auto bitset_original_nbits = bitset.get_original_nbits();
if (bitset_original_nbits == 0) {
bitset_original_nbits = sizeof(uint32_t) * 8;
}
RAFT_EXPECTS(
bitset.size() == n,
"IDSelectorBitmap is out of range for the given bitset");
"IDSelectorBitmap is out of range for the given bitset: %ld != %zu",
bitset.size(),
n);
auto stream = raft::resource::get_cuda_stream(res);
auto n_elements = (selector.n + 7) / 8;
auto d_bitmap =
raft::make_device_vector<uint8_t, uint32_t>(res, n_elements);
auto d_bitmap = raft::make_device_vector<uint8_t, int64_t>(res, n_elements);
auto d_bitmap_ptr = d_bitmap.data_handle();
raft::copy(
res,
d_bitmap.view(),
raft::make_host_vector_view<const uint8_t, uint32_t>(
raft::make_host_vector_view<const uint8_t, int64_t>(
selector.bitmap, n_elements));
const int threads_per_block = 256;
@@ -174,15 +175,15 @@ void convert_to_bitset_bitmap(
void convert_to_bitset_bruteforce(
raft::resources const& res,
const faiss::IDSelector& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset,
cuvs::core::bitset_view<uint32_t, int64_t> bitset,
int num_threads = 0) {
auto bitset_cpu =
raft::make_host_vector<uint32_t, uint32_t>(bitset.n_elements());
raft::make_host_vector<uint32_t, int64_t>(bitset.n_elements());
auto nbits = sizeof(uint32_t) * 8;
if (num_threads == 0)
num_threads = omp_get_max_threads();
#pragma omp parallel for num_threads(num_threads)
for (uint32_t i = 0; i < bitset.n_elements(); i++) {
for (int64_t i = 0; i < bitset.n_elements(); i++) {
uint32_t element = uint32_t{0};
for (uint32_t j = 0; j < nbits; j++) {
if (i * nbits + j < bitset.size() &&
@@ -199,7 +200,7 @@ void convert_to_bitset_bruteforce(
void convert_to_bitset(
faiss::gpu::GpuResources* res,
const faiss::IDSelector& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset,
cuvs::core::bitset_view<uint32_t, int64_t> bitset,
int num_threads) {
raft::device_resources& raft_handle = res->getRaftHandleCurrentDevice();
// If the selector is simple, we can use the specialized functions
faiss/gpu/utils/CuvsFilterConvert.h
+1 -1
View File
@@ -36,6 +36,6 @@ namespace faiss::gpu {
void convert_to_bitset(
faiss::gpu::GpuResources* res,
const faiss::IDSelector& selector,
cuvs::core::bitset_view<uint32_t, uint32_t> bitset,
cuvs::core::bitset_view<uint32_t, int64_t> bitset,
int num_threads = 0);
} // namespace faiss::gpu