Introduce sample filtering to IVFPQ index search by alexanderguzhva · Pull Request #1513 · rapidsai/raft

alexanderguzhva · 2023-05-14T15:23:02Z

A prototype that introduces a per-sample filtering for IVFPQ search. Please feel free to use it as a foundation for the future change, if appropriate, because the code is functional, but is not super clean-and-neat.

The diff introduces a template parameter called SampleFilterT. An instance is expected

to be SampleFilterT() constructible (which was mostly needed to define a default behavior in the form of SampleFilterT sample_filter=SampleFilterT(), see below)
to provide a inline __device__ bool operator(...) that returns true is a given sample is valid for being used against a given query in IVFPQ search

The default filter (that I set as a default one in certain facilities in the form of typename SampleFilterT = NoneSampleFilter in order not to modify way to many files) allows all samples to be used:

struct NoneSampleFilter {
  inline __device__ __host__ bool operator()(
    // query index
    const uint32_t query_ix,
    // the current inverted list index
    const uint32_t cluster_ix,
    // the index of the current sample inside the current inverted list
    const uint32_t sample_ix
  ) const {
    return true;
  }
};

Here __host__ is needed for a CPU-based testing only.
Also, I've provided an implementation of BitMaskSampleFilter that allows to filter samples based on a bit mask, as an example. The implementation was tested in the semi-production environment.

All the filter-related files were added to cpp/include/raft/neighbors/detail/sample_filter.cuh.

I did not change the default ivf_pq_search() method remains unchanged, but one more ivf_pq_search_with_filtering() method with an additional template argument SampleFilterT and one more input parameter was introduced.

template <typename T, typename IdxT, typename SampleFilterT>
void search_with_filtering(raft::device_resources const& handle,
            const raft::neighbors::ivf_pq::search_params& params,
            const index<IdxT>& idx,
            const T* queries,
            uint32_t n_queries,
            uint32_t k,
            IdxT* neighbors,
            float* distances,
            rmm::mr::device_memory_resource* mr = nullptr,
            SampleFilterT sample_filter = SampleFilterT());

All the current instantiations use NoneSampleFilter only.

I've used SampleFilterT sample_filter parameters passing instead of const SampleFilterT sample_filter in the function calls in order to be able to add some debugging facilities to a filter and with the hope that the compiler is smart enough to understand the de-facto constness if needed.

The filter does not take a computed distance score into account by design, thus the current implementation cannot have a distance threshold. This can be easily changed, if appropriate.

It is still questionable to me whether this filtering needs to be injected right inside the search kernel instead of doing post-processing, please let me know if you have any thoughts on the topic.

I'm happy to address the comments.

Thanks.

…Please feel free to use it as a foundation for the future change, if appropriate.

rapids-bot · 2023-05-14T15:23:09Z

Pull requests from external contributors require approval from a rapidsai organization member with write or admin permissions before CI can begin.

cjnolet · 2023-05-14T16:40:55Z

/ok to test

achirkin

Thanks for the PR! The change looks good and neat. Though I think a few things require further discussion.

cpp/include/raft/neighbors/detail/sample_filter.cuh

cpp/include/raft/neighbors/detail/ivf_pq_compute_similarity-inl.cuh

cpp/include/raft/neighbors/detail/sample_filter.cuh

alexanderguzhva · 2023-05-16T13:39:41Z

@achirkin I've made the changes to the PR, please feel free to take a look whenever you have time. Thanks!

achirkin · 2023-05-16T15:49:32Z

Thanks, @alexanderguzhva. In the meanwhile, could you please elaborate a bit on the need for the query_ix parameter in the filtering? My concern is that the extra kernel parameter (query offset) increases the register pressure even if we don't do any filtering.

alexanderguzhva · 2023-05-16T16:12:32Z

Thanks, @alexanderguzhva. In the meanwhile, could you please elaborate a bit on the need for the query_ix parameter in the filtering? My concern is that the extra kernel parameter (query offset) increases the register pressure even if we don't do any filtering.

Well, the use case I'm interested in is when every query uses its own set of allowed samples. This is the requirement from the interested party.

cjnolet · 2023-05-16T21:10:24Z

Thanks, @alexanderguzhva. In the meanwhile, could you please elaborate a bit on the need for the query_ix parameter in the filtering?

@achirkin other folks have also requested that we support the ability to pre-filter (before or during k-selection, as opposed to afterwards) based on a given set of indices.

cjnolet · 2023-05-16T22:02:04Z

/ok to test

cjnolet · 2023-05-17T13:14:45Z

Linking FAISS PR containing design that this is enabling: facebookresearch/faiss#2848

alexanderguzhva · 2023-05-17T20:38:18Z

Anything needs to be done on my end so far? Thanks

cjnolet · 2023-05-18T00:19:02Z

/ok to test

cjnolet · 2023-05-18T00:51:19Z

/ok to test

cjnolet · 2023-05-18T01:54:54Z

/ok to test

alexanderguzhva · 2023-05-18T01:55:58Z

/ok to test

@cjnolet , do you have any handy script for running a style check locally? Thanks!

cjnolet · 2023-05-18T01:57:46Z

@alexanderguzhva we have pre-commit set up on the repository so it should do all the clang-formatting for you upon committing to git. Here's the instructions to install / configure it: https://docs.rapids.ai/api/raft/nightly/contributing/#python-pre-commit-hooks.

alexanderguzhva · 2023-05-18T02:01:00Z

@cjnolet Thanks a lot!

cjnolet · 2023-05-18T02:05:49Z

/ok to test

cjnolet · 2023-05-18T02:14:33Z

cpp/include/raft/neighbors/detail/ivf_pq_search.cuh

 }

 /** See raft::spatial::knn::ivf_pq::search docs */
+<<<<<<< HEAD


I think there might be some unmerged bits here.

Yep, I'll fix that, thanks!

cjnolet · 2023-05-18T20:53:36Z

/ok to test

achirkin

I've compared the prim benchmarks of this PR (with NoneSampleFilter) against the current 23.06. There seem to be no performance hit whatsoever. The ncu stats, such as the number of registers stay the same as well, despite having more arguments in the signature. That's an LGTM from the performance side of things!

cpp/include/raft/neighbors/ivf_pq-inl.cuh

cjnolet · 2023-05-19T12:30:34Z

/ok to test

cjnolet · 2023-05-19T22:37:16Z

/merge

cjnolet · 2023-05-19T22:38:09Z

Thanks again for contributing this feature @alexanderguzhva!

…anup (#1541) The PR does the following: * Introduces `ivf_flat::search_with_filtering()` call in the same way the filtering was introduced to ivf_pq in #1513 * Moves `sample_filter.cuh` from `raft/neighbor/detail` to `raft/neighbor` * Moves `NoneSampleFilter` from `raft::neighbor::ivf_pq::detail` namespace to `raft::neighbor::filtering` namespace * Renames `NoneSampleFilter` to `NoneIvfSampleFilter` and template argument `SampleFilterT` to `IvfSampleFilterT` * Adds a missing `resource::get_workspace_resource(handle)` in `ivf_flat-inl.cuh` in a `search_with_filtering()` call (which was copied from `search()` call with the same problem) * Adds more comments in `ivf_pq-inl.h` * Some code cleanup in `ivf_pq-inl.h` Authors: - Alexander Guzhva (https://github.com/alexanderguzhva) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) URL: #1541

A prototype that introduces a per-sample filtering for IVFPQ search. Please feel free to use it as a foundation for the future change, if appropriate, because the code is functional, but is not super clean-and-neat. The diff introduces a template parameter called `SampleFilterT`. An instance is expected * to be `SampleFilterT()` constructible (which was mostly needed to define a default behavior in the form of `SampleFilterT sample_filter=SampleFilterT()`, see below) * to provide a `inline __device__ bool operator(...)` that returns `true` is a given sample is valid for being used against a given query in IVFPQ search The default filter (that I set as a default one in certain facilities in the form of `typename SampleFilterT = NoneSampleFilter` in order not to modify way to many files) allows all samples to be used: ``` struct NoneSampleFilter { inline __device__ __host__ bool operator()( // query index const uint32_t query_ix, // the current inverted list index const uint32_t cluster_ix, // the index of the current sample inside the current inverted list const uint32_t sample_ix ) const { return true; } }; ``` Here `__host__` is needed for a CPU-based testing only. Also, I've provided an implementation of `BitMaskSampleFilter` that allows to filter samples based on a bit mask, as an example. The implementation was tested in the semi-production environment. All the filter-related files were added to `cpp/include/raft/neighbors/detail/sample_filter.cuh`. I did not change the default `ivf_pq_search()` method remains unchanged, but one more `ivf_pq_search_with_filtering()` method with an additional template argument `SampleFilterT` and one more input parameter was introduced. ``` template <typename T, typename IdxT, typename SampleFilterT> void search_with_filtering(raft::device_resources const& handle, const raft::neighbors::ivf_pq::search_params& params, const index<IdxT>& idx, const T* queries, uint32_t n_queries, uint32_t k, IdxT* neighbors, float* distances, rmm::mr::device_memory_resource* mr = nullptr, SampleFilterT sample_filter = SampleFilterT()); ``` All the current instantiations use `NoneSampleFilter` only. I've used `SampleFilterT sample_filter` parameters passing instead of `const SampleFilterT sample_filter` in the function calls in order to be able to add some debugging facilities to a filter and with the hope that the compiler is smart enough to understand the de-facto constness if needed. The filter does not take a computed distance score into account by design, thus the current implementation cannot have a distance threshold. This can be easily changed, if appropriate. It is still questionable to me whether this filtering needs to be injected right inside the search kernel instead of doing post-processing, please let me know if you have any thoughts on the topic. I'm happy to address the comments. Thanks. Authors: - Alexander Guzhva (https://github.com/alexanderguzhva) - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) URL: rapidsai#1513

…anup (rapidsai#1541) The PR does the following: * Introduces `ivf_flat::search_with_filtering()` call in the same way the filtering was introduced to ivf_pq in rapidsai#1513 * Moves `sample_filter.cuh` from `raft/neighbor/detail` to `raft/neighbor` * Moves `NoneSampleFilter` from `raft::neighbor::ivf_pq::detail` namespace to `raft::neighbor::filtering` namespace * Renames `NoneSampleFilter` to `NoneIvfSampleFilter` and template argument `SampleFilterT` to `IvfSampleFilterT` * Adds a missing `resource::get_workspace_resource(handle)` in `ivf_flat-inl.cuh` in a `search_with_filtering()` call (which was copied from `search()` call with the same problem) * Adds more comments in `ivf_pq-inl.h` * Some code cleanup in `ivf_pq-inl.h` Authors: - Alexander Guzhva (https://github.com/alexanderguzhva) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) URL: rapidsai#1541

A prototype that introduces a per-sample filtering for IVFPQ search. …

da7538c

…Please feel free to use it as a foundation for the future change, if appropriate.

alexanderguzhva requested a review from a team as a code owner May 14, 2023 15:23

github-actions bot added the cpp label May 14, 2023

cjnolet added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels May 14, 2023

achirkin requested changes May 15, 2023

View reviewed changes

address PR comments

24e4bdf

cjnolet assigned alexanderguzhva May 15, 2023

Merge branch 'branch-23.06' into branch-23.06

0d8a2a1

fix code style

f6b0339

alexanderguzhva added 2 commits May 18, 2023 00:31

fix code style

90893cd

fix code style again

9653cc2

alexanderguzhva added 2 commits May 17, 2023 21:35

fix a missing NoneSampleFilter in ivf_pq_compute_similarity-ext.cuh

197b9ab

upstream rapidsai/raft:branch23.06

53a4ef8

fix code style

0aa4d26

cjnolet reviewed May 18, 2023

View reviewed changes

cjnolet mentioned this pull request May 18, 2023

[FEA] Pre-filtering of nearest neighbors #1528

Open

5 tasks

fix merge errors and linter warnings

c150b2c

Merge remote-tracking branch 'upstream/branch-23.06' into branch-23.06

0a80332

achirkin approved these changes May 19, 2023

View reviewed changes

cpp/include/raft/neighbors/ivf_pq-inl.cuh Show resolved Hide resolved

Merge branch 'branch-23.06' into branch-23.06

cda0371

cjnolet approved these changes May 19, 2023

View reviewed changes

rapids-bot bot merged commit cdf107b into rapidsai:branch-23.06 May 19, 2023

alexanderguzhva mentioned this pull request May 20, 2023

Add sample filtering for ivf_flat. Filtering code refactoring and cleanup #1541

Merged

Conversation

alexanderguzhva commented May 14, 2023

Uh oh!

rapids-bot bot commented May 14, 2023

Uh oh!

cjnolet commented May 14, 2023

Uh oh!

achirkin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alexanderguzhva commented May 16, 2023

Uh oh!

achirkin commented May 16, 2023

Uh oh!

alexanderguzhva commented May 16, 2023

Uh oh!

cjnolet commented May 16, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cjnolet commented May 16, 2023

Uh oh!

cjnolet commented May 17, 2023

Uh oh!

alexanderguzhva commented May 17, 2023

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

alexanderguzhva commented May 18, 2023

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

alexanderguzhva commented May 18, 2023

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

cjnolet May 18, 2023

Choose a reason for hiding this comment

Uh oh!

alexanderguzhva May 18, 2023

Choose a reason for hiding this comment

Uh oh!

cjnolet commented May 18, 2023

Uh oh!

achirkin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cjnolet commented May 19, 2023

Uh oh!

cjnolet commented May 19, 2023

Uh oh!

cjnolet commented May 19, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cjnolet commented May 16, 2023 •

edited

Loading