Add explicit instantiations for IVF-PQ search kernels used in tests

[tfeher]

Compilation of IVF-PQ search kernels can be time consuming. In libraft.so the compilation is done in parallel for kernels without filtering and with int64_t index type.

We have test with uint32_t index type as well as tests for bitset_filter with both 32 and 64 bit index types. This PR adds explicit template instantiations for the test. This way we avoid repeated compilation of the kernels with filter and this also enables parallel compilation of the compute_similarity kernel for different template types. The kernels with these additional type parameters are not added to libraft.so, only linked together with the test executable.

Note that this PR does not increase the number of compiled kernels, but it enables to compile them in parallel.

[tfeher]

• By default ivf_pq::search calls search_with_filtering using a none_ivf_sample_filter.
• Internally we wrap the filter into a filter_adapter

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

  Lines 654 to 655 in 054731d

  	auto filter_adapter = raft::neighbors::filtering::ivf_to_sample_filter(
  	index.inds_ptrs().data_handle(), sample_filter);

• The filter adapter is dependent on the index type:

  raft/cpp/include/raft/neighbors/sample_filter_types.hpp

  Lines 68 to 76 in 054731d

  	template <typename index_t, typename filter_t>
  	struct ivf_to_sample_filter {
  	const index_t* const* inds_ptrs_;
  	const filter_t next_filter_;
  	ivf_to_sample_filter(const index_t* const* inds_ptrs, const filter_t next_filter)
  	: inds_ptrs_{inds_ptrs}, next_filter_{next_filter}
  	{
  	}

• The filter type is a template argument of compute_similarity_kernel.

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

  Lines 264 to 271 in 054731d

  	template <typename OutT,
  	typename LutT,
  	typename IvfSampleFilterT,
  	uint32_t PqBits,
  	int Capacity,
  	bool PrecompBaseDiff,
  	bool EnableSMemLut>
  	RAFT_KERNEL compute_similarity_kernel(uint32_t dim,

• The compute_similarity_kernel does not directly depend on IdxT. The neighbor indices are mapped to the requested index type in post processing. But IvfSampleFilterT's dependence on IdxT makes compute_similarity_kernel also dependent on the index type.
• This leads to the unfortunate situation that we need to instantiate the search kernels for uint32_t as well as int64_t, even though IdxT should not matter when none_ivf_sample_filter is used.
  https://github.com/rapidsai/raft/pull/2212/files#diff-d5a033b62a4c4c60c86931d10ba5f3b8b0a6e805342da7050829c7717c2305ccR406-R412

What is the best way to avoid this extra compilation @lowener and @achirkin?

One idea would be to cast away the index type if none_ivf_sample_filter is selected while we wrap it into the adapter. Could we have a way where the filter adapter is not templated on IdxT?

[cjnolet]

One idea would be to cast away the index type if none_ivf_sample_filter is selected while we wrap it into the adapter. Could we have a way where the filter adapter is not templated on IdxT?

We had a discussion a little while back of just having the search_with_filter() accept a cuda::std::function as an argument (thus removing the template for the filter type and turning it into a runtime argument). I wonder if it's finally time to make that change in the 24.04 version.

[tfeher]

I wonder if it's finally time to make that change in the 24.04 version.

As discussed offline, I recommend to push further changes to the filter function to 24.06, and keep this PR in its current, non-breaking form. I have added a reference to the tracking issue for filters #1738

[tfeher]

Overview of the changes in compile time. Test and bench with filtering plus test with uint32_t (highlighted in bold) is now compiled in parallel

file	build time before	after PR (mm:ss)
ann_ivf_pq/test_float_uint32_t.cu	50:15	3:13
ann_ivf_pq/test_filter_float_int64_t.cu	26:52	1:30
ann_ivf_pq/test_filter_int8_t_int64_t.cu	26:48	1:50
bench/ivf_pq_filter_float_int64_t.cu	21:05	1:48
ann_ivf_pq/test_int8_t_int64_t.cu	3:17	3:14
ann_ivf_pq/test_float_int64_t.cu	3:14	4:14
ann_ivf_pq/test_uint8_t_int64_t.cu	3:11	4:11
bench/ivf_pq_int8_t_int64_t.cu	1:55	1:45
bench/ivf_pq_uint8_t_int64_t.cu	1:53	1:39
bench/ivf_pq_float_int64_t.cu	1:53	1:56
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_true_bitset32	0 (new)	6:06
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_true_filt32	0 (new)	6:03
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_true_bitset64	0 (new)	5:37
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_false_filt32	0 (new)	5:15
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_false_bitset32	0 (new)	5:15
ann_ivf_pq/ivf_pq_compute_similarity_half_half_bitset32	0 (new)	5:06
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_true_bitset32	0 (new)	5:04
ann_ivf_pq/ivf_pq_compute_similarity_half_fp8_false_bitset64	0 (new)	5:03
ann_ivf_pq/ivf_pq_compute_similarity_half_half_filt32	0 (new)	5:03
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_true_filt32	0 (new)	5:02
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_true_bitset64	0 (new)	4:51
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_false_bitset32	0 (new)	4:46
ann_ivf_pq/ivf_pq_compute_similarity_half_half_bitset64	0 (new)	4:44
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_false_filt32	0 (new)	4:44
ann_ivf_pq/ivf_pq_compute_similarity_float_fp8_false_bitset64	0 (new)	4:42
ann_ivf_pq/ivf_pq_compute_similarity_float_half_bitset32	0 (new)	4:37
ann_ivf_pq/ivf_pq_compute_similarity_float_half_filt32	0 (new)	4:37
ann_ivf_pq/ivf_pq_compute_similarity_float_half_bitset64	0 (new)	4:35
ann_ivf_pq/ivf_pq_compute_similarity_float_float_bitset32	0 (new)	4:34
ann_ivf_pq/ivf_pq_compute_similarity_float_float_filt32	0 (new)	4:31
ann_ivf_pq/ivf_pq_compute_similarity_float_float_bitset64	0 (new)	4:28

[achirkin]

Thanks, @tfeher, for the PR! Looks good to me. A small nitpick about the internal indexing type below.

[benfred]

lgtm!

[tfeher]

/merge