fix knn_merge_parts for inner product

[qwertyforce]

By default, when merging the results of sharded multigpu index, knn_merge_parts keeps K smallest values. But inner product is not a distance, it is a measure of similarity. Therefore results are wrong for IP.
In this PR we are adding an additional overload, that receives an argument select_min, we keep backward compatibility and add a new test

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[divyegala]

Hi @qwertyforce, thank you for the PR! I think this change will severely affect the binary size of cuVS, can you please provide a before-after measurement? Alternatively, we could run an element-wise operation to negate the inner product distances before running the merge kernel.

[qwertyforce]

Hi! i ran ./build.sh libcuvs --allgpuarch --no-nvtx -n (GCC 13, CUDA 12.8)
main vs patched
libcuvs.so increased from 504MB to 512MB
knn_merge_parts.cu.o increased from 8MB to 16MB

[divyegala]

Thanks @qwertyforce , in that case can we please pursue the alternative of running a negation for inner product?

[qwertyforce]

Yes, will try to implement and benchmark it

[qwertyforce]

made a negation version, doesnt affect ivf search performance in multigpu index
similar to

cuvs/cpp/src/neighbors/detail/tiered_index.cuh

Lines 244 to 269 in 3937a7c

	if (!distance::is_min_close(build_params.metric)) {
	// knn_merge_parts doesn't currently support InnerProduct distances etc
	// instead negate here and then undo after
	raft::linalg::map(res,
	temp_distances.view(),
	raft::mul_const_op<value_type>(-1),
	raft::make_const_mdspan(temp_distances.view()));
	}
	// merge results from ann_index/bfknn together, translating the bfknn ids
	auto stream = raft::resource::get_cuda_stream(res);
	int64_t host_translations[2] = {0, static_cast<int64_t>(ann_rows())};
	auto device_translations = raft::make_device_vector<int64_t>(res, 2);
	raft::copy(device_translations.data_handle(), host_translations, 2, stream);
	knn_merge_parts(res,
	temp_distances.view(),
	temp_neighbors.view(),
	distances,
	neighbors,
	device_translations.view());
	if (!distance::is_min_close(build_params.metric)) {
	raft::linalg::map(
	res, distances, raft::mul_const_op<value_type>(-1), raft::make_const_mdspan(distances));
	}

also, as i understand, it is possible that other instances where knn_merge_parts is used could also be affected, for example (should probably be resolved in another pr)

cuvs/cpp/src/neighbors/detail/sparse_knn.cuh

Lines 317 to 375 in cd4bbab

	void merge_batches(csr_batcher_t<value_idx, value_t>& idx_batcher,
	csr_batcher_t<value_idx, value_t>& query_batcher,
	value_t* merge_buffer_dists,
	value_idx* merge_buffer_indices,
	value_t* out_dists,
	value_idx* out_indices)
	{
	// build translation buffer to shift resulting indices by the batch
	std::vector<value_idx> id_ranges;
	id_ranges.push_back(0);
	id_ranges.push_back(idx_batcher.batch_start());
	rmm::device_uvector<value_idx> trans(id_ranges.size(), raft::resource::get_cuda_stream(handle));
	raft::copy(handle,
	raft::make_device_vector_view(trans.data(), id_ranges.size()),
	raft::make_host_vector_view(id_ranges.data(), id_ranges.size()));
	// combine merge buffers only if there's more than 1 partition to combine
	auto rows = query_batcher.batch_rows();
	knn_merge_parts(
	handle,
	raft::make_device_matrix_view<const value_t, int64_t>(merge_buffer_dists, rows, 2 * k),
	raft::make_device_matrix_view<const value_idx, int64_t>(merge_buffer_indices, rows, 2 * k),
	raft::make_device_matrix_view<value_t, int64_t>(out_dists, rows, k),
	raft::make_device_matrix_view<value_idx, int64_t>(out_indices, rows, k),
	raft::make_device_vector_view<value_idx, int64_t>(trans.data(), id_ranges.size()));
	}
	void perform_k_selection(csr_batcher_t<value_idx, value_t> idx_batcher,
	csr_batcher_t<value_idx, value_t> query_batcher,
	value_t* batch_dists,
	value_idx* batch_indices,
	value_t* out_dists,
	value_idx* out_indices)
	{
	// populate batch indices array
	value_idx batch_rows = query_batcher.batch_rows(), batch_cols = idx_batcher.batch_rows();
	// build translation buffer to shift resulting indices by the batch
	std::vector<value_idx> id_ranges;
	id_ranges.push_back(0);
	id_ranges.push_back(idx_batcher.batch_start());
	// in the case where the number of idx rows in the batch is < k, we
	// want to adjust k.
	value_idx n_neighbors = std::min(static_cast<value_idx>(k), batch_cols);
	bool ascending = cuvs::distance::is_min_close(metric);
	// kernel to slice first (min) k cols and copy into batched merge buffer
	cuvs::selection::select_k(
	handle,
	raft::make_device_matrix_view<const value_t, int64_t>(batch_dists, batch_rows, batch_cols),
	raft::make_device_matrix_view<const value_idx, int64_t>(
	batch_indices, batch_rows, batch_cols),
	raft::make_device_matrix_view<value_t, int64_t>(out_dists, batch_rows, n_neighbors),
	raft::make_device_matrix_view<value_idx, int64_t>(out_indices, batch_rows, n_neighbors),
	ascending,
	true);

[divyegala]

@viclafargue could you please take a look at whether this PR needs follow-on work, and if yes, capture it an issue?

[divyegala]

Thanks for the contribution!

[divyegala]

/ok to test bb6d2df

[viclafargue]

Great catch! Thank you for identifying the issue and following through with a PR to fix it.

@divyegala There indeed is an other independently existing issue in the direct merge path that would deserve some work. I will directly open up a PR to fix this.

EDIT: Here is the PR : #2441.

[viclafargue]

/merge