From dbfa7ddb3c0008852d7a2c0c672af19aaf3f679d Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Tue, 13 Feb 2024 14:03:36 +0100 Subject: [PATCH 1/2] Fix ANN bench ground truth generation for k>1024 --- cpp/include/raft/neighbors/detail/knn_merge_parts.cuh | 3 +++ .../src/raft-ann-bench/generate_groundtruth/__main__.py | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/knn_merge_parts.cuh b/cpp/include/raft/neighbors/detail/knn_merge_parts.cuh index c8ff03741c..bc64fc9181 100644 --- a/cpp/include/raft/neighbors/detail/knn_merge_parts.cuh +++ b/cpp/include/raft/neighbors/detail/knn_merge_parts.cuh @@ -16,6 +16,7 @@ #pragma once +#include #include #include @@ -168,5 +169,7 @@ inline void knn_merge_parts(const value_t* inK, else if (k <= 1024) knn_merge_parts_impl( inK, inV, outK, outV, n_samples, n_parts, k, stream, translations); + else + THROW("Unimplemented for k=%d, knn_merge_parts works for k<=1024", k); } } // namespace raft::neighbors::detail diff --git a/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py index f4d97edea5..2250618f58 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py @@ -67,12 +67,13 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"): queries, k, metric=metric, - handle=handle, - global_id_offset=i, # shift neighbor index by offset i + handle=handle ) handle.sync() D, Ind = cp.asarray(D), cp.asarray(Ind) + Ind += i # shift neighbor index by offset i + if distances is None: distances = D indices = Ind From 485460932fe14c7c522630343015c248b9bd1ded Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Tue, 13 Feb 2024 14:35:57 +0100 Subject: [PATCH 2/2] Fix style --- .../raft-ann-bench/generate_groundtruth/__main__.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py index 2250618f58..a5ebb76635 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/generate_groundtruth/__main__.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # -# Copyright (c) 2023, NVIDIA CORPORATION. +# Copyright (c) 2023-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -62,18 +62,12 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"): X = cp.asarray(dataset[i : i + n_batch, :], cp.float32) - D, Ind = knn( - X, - queries, - k, - metric=metric, - handle=handle - ) + D, Ind = knn(X, queries, k, metric=metric, handle=handle) handle.sync() D, Ind = cp.asarray(D), cp.asarray(Ind) Ind += i # shift neighbor index by offset i - + if distances is None: distances = D indices = Ind