[FEA] Support non-fused k-selection in IVF-flat

[lowener]

raft/cpp/include/raft/matrix/detail/select_warpsort.cuh

Line 100 in 99655ea

static constexpr int kMaxCapacity = 256;

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

Lines 1003 to 1008 in 99655ea

	// NB: this is the limitation of the warpsort structures that use a huge number of
	// registers (used in the main kernel here).
	RAFT_EXPECTS(capacity == Capacity,
	"Capacity must be power-of-two not bigger than the maximum allowed size "
	"matrix::detail::select::warpsort::kMaxCapacity (%d).",
	matrix::detail::select::warpsort::kMaxCapacity);

Providing a fallback when k>256 would be useful.

[cjnolet]

Linking this to #711, converting this to a feature request, and tagging @achirkin to provide details Im likely to miss- the problem here is that ivf-pq supports both fused and non-fused k-selection (and so is able to support very large settings of k) while ivf-flat on supports the fused k-selection.

The need here is to implement non-fused k-selection in ivf-flat so that we can support larger values of k.

[achirkin]

As @cjnolet aptly specified in the new title, the issue is that IVF-flat implementation relies on select::warpsort to fuse the k-selection with computing the distance. The warpsort primitives store the top-k values per-warp in registers for efficiency, and thus k cannot grow unbound. The only workaround I see is to implement a non-fused version of IVF-flat, that would use the host-side matrix::select_k (which falls back to the radix-based implementation for large k). This would be the same approach as the currently implemented IVF-PQ.

Following the IVF-PQ example, I think the scope of the work can (mostly) be limited to a single file neighbors/detail/ivf_flat_interleaved_scan-inl.cuh:

1. interleaved_scan_kernel: add a new codepath re-using template parameter Capacity == 0 we may need to factor-out and reuse ivf-pq's dummy_block_sort_t for this.
2. NB: In the non-fused mode, each processed (query, probe) produces the output of different size (number of rows in a cluster). We'd need to tackle this either by padding data or calculating the sizes. In ivf-pq, we estimate max_samples - maximum possible cumulative size of clusters probed for one query - to allocate a big enough intermediate buffer for the found distances. The results are packed densely for probes within one query, but padded with dummy values between queries.
3. For (2) we'd need something similar to calc_chunk_indices
4. In the non-fused version of the kernel we don't store the indices to save the memory bandwidth. Thus the indices need to be reconstructed in an extra post-processing step (see ivf-pq's postprocess_neighbors)
5. As an extra to the previous step, we can improve compile times by fixing the IdxT = uint32_t of the interleaved_scan_kernel (Make all codepaths output the sample indices and convert them to the db indices in the post-processing step).
6. Adjust the select_interleaved_scan_kernel to include the Capacity == 0 case (a minimal change).
7. Since we already have some batching inside the launch_kernel, we could allocate the extra intermediate buffer storing the distances (max_samples * kMaxGridY) outside that loop and reuse it. Then the matrix::select_k would go at the end of each iteration.