Instantiate uint64_t IndexType for CAGRA search kernels

[divyegala]

These instantiations are required in facebookresearch/faiss#3084 to link to precompiled CAGRA binaries from libraft.so

[tfeher]

/rerun tests

[achirkin]

Looks like the *.cu files are missing from the PR?

[tfeher]

Thanks Divye for this PR! Currently INDEX_T is used for three different purpose:

1. Output neighbor index type
2. Graph index type
3. Index type used in search kernels to store visited / candidate nodes.

The goal of this PR is to change 1. to int64_t to facilitate FAISS integration. But the current PR effectively changes all 3 types, and that is not ideal. We have the following problems:

2. Graph size. A graph for 100M vectors and 32 degree is already 12 GiB. This would double, reducing the data size that we can handle in GPU memory.

3. a) Binary size: currently CAGRA search kernels add around 50 MiB to libraft.so. I expect that this would double, and also the build time would correspondingly increase (although it can be build in parallel, so this might be less concern).

   b) Performance: the kernels are optimized for 32bit indexing. We should benchmark and possibly adjust the kernel configs for 64bit index type.

Considering these issues I would recommend to separate these concerns, and modify only the output index type (1.)

This can be easily done by introducing an extra output index template parameter, and an elementwise op to cast the index type if necessary. This is already used in our benchmark, so our benchmark times already include this overhead:

raft/cpp/bench/ann/src/raft/raft_cagra_wrapper.h

Lines 283 to 289 in 7342980

	if constexpr (!std::is_same_v<IdxT, size_t>) {
	raft::linalg::unaryOp(neighbors,
	neighbors_IdxT,
	batch_size * k,
	raft::cast_op<size_t>(),
	raft::resource::get_cuda_stream(handle_));
	}

If we have any performance consideration with additional kernel call, then I would recommend to add a follow-up task to improve that. One could add the cast when

• (for single-cta algo) the neighbor indices are written back to global memory by the search kernel,
• (for multi-cta algo) top-k results are selected from intermediate results
• (for multi-kernel algo) copy results to final buffer.

While changing single-cta kernel, one should watch out for binary size: we can do a runtime type dispatch without introducing an additional template parameter.

[divyegala]

@tfeher thanks for the detailed writeup! I agree with you that we should limit the changes just to the OutputType. Would you agree that in this case where we have a lot of considerations to figure out, we should just add a copy in FAISS itself to cast the output from uint32_t to int64_t for the first iteration, then follow up with updates to see if we can natively support OutputType=int64_t with acceptable changes to binary size and overheads?

[tfeher]

add a copy in FAISS itself to cast the output from uint32_t to int64_t for the first iteration, then follow up with updates to see if we can natively support OutputType=int64_t

Depends. If you anyways have to map the neighbor idx data in the faiss integration (e.g. marking invalid values according to faiss expectation), then it sounds fine to me to do the mapping in faiss.

Otherwise I would prefer to do the cast in cagra::search, because then we need to modify only one library when we decide to further optimize this.

[divyegala]

Depends. If you anyways have to map the neighbor idx data in the faiss integration (e.g. marking invalid values according to faiss expectation), then it sounds fine to me to do the mapping in faiss.

@tfeher I'm not sure what this means. When would values be invalid in faiss?