IVF-PQ: manipulating individual lists #1298

achirkin · 2023-02-23T12:55:33Z

Add public functions for reading and writing into individual ivf-pq lists (clusters), in the input space (reconstructed data) and in flat PQ codes.

Partially solves (IVF-PQ) #1205

achirkin · 2023-03-07T06:08:47Z

cpp/include/raft/neighbors/ivf_pq.cuh

+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::device_resources const& res,
+                           const index<IdxT>& index,
+                           device_matrix_view<T, uint32_t, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t n_skip)


@cjnolet , could you please check if this interface is suitable for integration with FAISS?

A couple important points that the FAISS developers have pointed out:

We need to be sure that we don't assume indexes are drawn from a strictly monotonically increasing set.

Following from 1, what's been asked for is a reconstruct_batch API which accepts a collection of ids (which are not necessarily contiguous) and reconstructs those vectors.

From what I can tell, the API you are proposing here doesn't implement the above, so I would say it's not suitable for the FAISS integration.

A quick clarification: the ids we should pass are (A) the same indices we pass during construction or (B) internal ids of the ivf-pq index?
If (B) then should these be a pair (label, collection of in-cluster-ids), or a collection of (label, in-cluster-id)? The former, perhaps can be slightly faster.
If (A) then I hardly can imagine an efficient function that can do this: we don't sort or organize the input indices in any way, so a full search in the database needs to be done for every index. Perhaps we can split the task in two: search indices -> then call interface (B)?

Also don't you need something like this interface to be able to convert our GPU index to their CPU index type (more-or-less efficiently, converting one cluster at a time)?

Also don't you need something like this interface to be able to convert our GPU index to their CPU index type (more-or-less efficiently, converting one cluster at a time)?

There's a couple different asks here. We have an ask to be able to convert from one interleaved (and potentially non-interleaved) encoded format to another, which is required for being able to convert between cpu and gpu indices. That ask is for us to have functions available that can be used by their CodePacker API.

In addition to this, there's the ask to be able to reconstruct a batch of vectors given a set of indices, which might not be in order. This is the corresponding issue for that ask (which is directly from Jeff Johnson). The indices passed into the reconstruct_batch would also coincide w/ the indices which would be specified when invoking extend here.

That says to the world "We will not change these APIs, nor the way they function, without deprecation first".

Oh, now I see what you mean. I agree, we should export detail::pack_list_data and detail::unpack_list_data into a public namespace, to make sure we don't change their signature and break faiss integration by accident. I'm also not opposed to the idea of moving the construct/reconstruct and other similar functions to the ivf_pq::helpers public namespace.
The only thing that bothers me, is that the overloads we need for the codepacker go really against the current design of the public api. Maybe, as an exception, we export these two into something like ivf_pq::codepacker to discourage other users from using them? Or even implement the whole codepacker api in raft in that namespace?

to discourage other users from using them?

I really don't think this is a good way to approach any public API. Either the API is public and users can use it, or it's not and they should touch that code at their own risk. Trying to force this middle ground for some users and not others is not sustainable and I'm not sure why there's a problem just exposing it outright.

I also don't agree that these functions should not just be part of the public API- if users need the ability to call a function, it needs to wrapped into the public API so it can be clearly documented, not just exported.

We had tried the exporting initially when we first started exposing formal public APIs and that was a mistake. Anything the user should be interacting directly with should be available and clearly documented (which for most almost all of the functions in RAFT just means wrapped).

I'm not opposed to using a nested "CodePacker" namespace but I'd also prefer that it is also nested inside helpers, so we be consistent there wrt the standard API namespaces and the "extra" functionality.

I have added a public API for pack_list_data.

Note, for the codepacker, we'd probably need to expose this overload

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

Lines 575 to 580 in 926f510

inline void unpack_list_data(

device_matrix_view<uint8_t, uint32_t, row_major> codes,

device_mdspan<const uint8_t, list_spec<uint32_t, uint32_t>::list_extents, row_major> list_data,

std::variant<uint32_t, const uint32_t*> offset_or_indices,

uint32_t pq_bits,

rmm::cuda_stream_view stream)

I have missed this message before. I have exposed the pack/unpack versions that take list_data directly under ivf_pq::helpers::codepacker namespace.

In our implementation we still need to provide the pq_bits parameter. Otherwise it is very similar to the CodePacker API.

One question: what does block refer to in CodePacker? In our API it refers to a set of interleaved vectors, which can be the data for a whole cluster, or less.

tfeher

Thanks Artem for the PR! As far as I understand we are still waiting for additional helpers to enable looking up vectors using the original index. Until that is ready, I have looked at the implementation details. It looks good so far, below you will find my comments.

cpp/include/raft/neighbors/ivf_pq.cuh

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

…reconstruct

cjnolet

I think this generally looks good, though I'll need to go back over the codepacker API just to make sure we'll be able to fill it in with these functions. Just a few other (mostly minor) API-ish things from my side.

Note, although the public api provides now a way to read raw flat PQ codes, it's not suitable for the CodePacker, because it's not raw pointers. However, one of the corresponding functions in detail namespace hopefully should work for this purpose (detail::pack_list_data and detail::unpack_list_data).

I think we should use raft::span here to expose these in raft::neighbors::ivf_pq::helpers, especially since we know the block will be of size block_size and the flat_codes will be size nvec * code_size.

cpp/include/raft/neighbors/ivf_pq.cuh

achirkin · 2023-03-28T08:11:51Z

I'm not sure, what do you mean by using raft::span for the codepacker? Probably, this my comment was a bit misleading. I think the problem with codepacker integration is that it takes the pointers to arbitrary locations within the index / lists as arguments. I don't think we would want to expose such a matching public function, because it goes against the hiding implementation details logic. Instead, I chose to reorganize the internal code, so that detail::pack_list_data and detail::unpack_list_data matched the codepacker interface and didn't take the index itself as an argument (nb: it's still mdspans, not raw pointers).

tfeher · 2023-03-31T15:39:07Z

@cjnolet Since Artem is on leave, I am addressing the open issues. Will push update later today.

tfeher

I have addressed the open issues.

cpp/include/raft/neighbors/ivf_pq.cuh

tfeher · 2023-03-31T20:17:03Z

cpp/include/raft/neighbors/ivf_pq.cuh

+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::device_resources const& res,
+                           const index<IdxT>& index,
+                           device_matrix_view<T, uint32_t, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t n_skip)


I have added a public API for pack_list_data.

cpp/test/neighbors/ann_ivf_pq.cuh

cjnolet

LGTM. Thanks Tamas!

cjnolet · 2023-04-17T20:06:28Z

/merge

Add public functions for reading and writing into individual ivf-pq lists (clusters), in the input space (reconstructed data) and in flat PQ codes. Partially solves (IVF-PQ) rapidsai#1205 Authors: - Artem M. Chirkin (https://github.com/achirkin) - Corey J. Nolet (https://github.com/cjnolet) - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: rapidsai#1298

Add a method to reconstruct the compressed index data

d421375

achirkin added feature request New feature or request non-breaking Non-breaking change 2 - In Progress Currenty a work in progress labels Feb 23, 2023

achirkin requested a review from a team as a code owner February 23, 2023 12:55

github-actions bot added the cpp label Feb 23, 2023

achirkin mentioned this pull request Feb 23, 2023

IVF-Flat reconstruction #1270

Open

achirkin added 3 commits February 23, 2023 13:59

Fix a typo in the docs

c7b5574

Relax the constraints a bit

087919f

Add the public interface

b5e9844

achirkin added 3 - Ready for Review and removed 2 - In Progress Currenty a work in progress labels Feb 23, 2023

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

ea81c1e

achirkin self-assigned this Feb 27, 2023

achirkin added 3 commits February 28, 2023 07:09

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

6567f5c

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

3ccd48f

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

fdb8a6a

achirkin requested review from tfeher and cjnolet March 7, 2023 06:05

achirkin commented Mar 7, 2023

View reviewed changes

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

56ffc84

tfeher reviewed Mar 13, 2023

View reviewed changes

cpp/include/raft/neighbors/ivf_pq.cuh Outdated Show resolved Hide resolved

cpp/include/raft/neighbors/detail/ivf_pq_build.cuh Outdated Show resolved Hide resolved

achirkin added 4 commits March 15, 2023 10:40

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

ebd1b1c

Fix the merge errors

2773a86

Add an option to reconstruct cluster data by in-cluster indices

82be40c

Update the docs for the new function

5fc1c53

achirkin added 2 - In Progress Currenty a work in progress and removed 3 - Ready for Review labels Mar 15, 2023

achirkin added 2 commits March 15, 2023 19:25

Detach reconstruction logic from the data reading logic

ddc4466

Merge remote-tracking branch 'rapidsai/branch-23.04' into fea-ivf-pq-…

56eb716

…reconstruct

achirkin added the 3 - Ready for Review label Mar 17, 2023

achirkin and others added 10 commits March 17, 2023 13:37

Merge remote-tracking branch 'rapidsai/branch-23.04' into fea-ivf-pq-…

2625666

…reconstruct

Merge remote-tracking branch 'rapidsai/branch-23.04' into fea-ivf-pq-…

172673b

…reconstruct

Relax the test criterion eps a little bit

1504027

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

1326ac7

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

08cfacb

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

a399023

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

d935568

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

d1dd238

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

5514d7d

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

a317622

cjnolet requested changes Mar 27, 2023

View reviewed changes

cpp/include/raft/neighbors/ivf_pq.cuh Outdated Show resolved Hide resolved

cpp/include/raft/neighbors/ivf_pq.cuh Outdated Show resolved Hide resolved

cpp/include/raft/neighbors/ivf_pq.cuh Outdated Show resolved Hide resolved

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

9367936

tfeher added 3 commits March 31, 2023 21:08

Move ivf_pq helpers to separate file and namespace

09f1a0d

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

b7e811a

Add public API for pack_list_data

926f510

tfeher reviewed Mar 31, 2023

View reviewed changes

tfeher added 4 commits April 2, 2023 21:39

Increase tolarance for vector reconstruction test

8e311d9

Merge branch 'branch-23.04' into fea-ivf-pq-reconstruct

158189c

Correct number of list elements to compare

280c040

Add ivf_pq::helpers::codepacker::pack / unpack

a7e5419

cjnolet changed the base branch from branch-23.04 to branch-23.06 April 12, 2023 17:50

Merge branch 'branch-23.06' into fea-ivf-pq-reconstruct

09abadb

cjnolet approved these changes Apr 14, 2023

View reviewed changes

Merge branch 'branch-23.06' into fea-ivf-pq-reconstruct

923c3b5

rapids-bot bot merged commit ba207a0 into rapidsai:branch-23.06 Apr 17, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

IVF-PQ: manipulating individual lists #1298

IVF-PQ: manipulating individual lists #1298

achirkin commented Feb 23, 2023 •

edited

Loading

achirkin Mar 7, 2023

cjnolet Mar 7, 2023

achirkin Mar 7, 2023

achirkin Mar 7, 2023

cjnolet Mar 7, 2023

achirkin Mar 29, 2023

cjnolet Mar 30, 2023

tfeher Mar 31, 2023

achirkin Mar 31, 2023

tfeher Apr 5, 2023 •

edited

Loading

tfeher left a comment

cjnolet left a comment •

edited

Loading

achirkin commented Mar 28, 2023

tfeher commented Mar 31, 2023

tfeher left a comment

tfeher Mar 31, 2023

cjnolet left a comment

cjnolet commented Apr 17, 2023

	inline void unpack_list_data(
	device_matrix_view<uint8_t, uint32_t, row_major> codes,
	device_mdspan<const uint8_t, list_spec<uint32_t, uint32_t>::list_extents, row_major> list_data,
	std::variant<uint32_t, const uint32_t*> offset_or_indices,
	uint32_t pq_bits,
	rmm::cuda_stream_view stream)

IVF-PQ: manipulating individual lists #1298

IVF-PQ: manipulating individual lists #1298

Conversation

achirkin commented Feb 23, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tfeher Apr 5, 2023 • edited Loading

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

cjnolet left a comment • edited Loading

Choose a reason for hiding this comment

achirkin commented Mar 28, 2023

tfeher commented Mar 31, 2023

tfeher left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cjnolet left a comment

Choose a reason for hiding this comment

cjnolet commented Apr 17, 2023

achirkin commented Feb 23, 2023 •

edited

Loading

tfeher Apr 5, 2023 •

edited

Loading

cjnolet left a comment •

edited

Loading