Fix ANN bench ground truth generation for k>1024

rapidsai · Feb 13, 2024 · dbfa7dd · dbfa7dd
1 parent c9574d7
commit dbfa7dd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git 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 <raft/core/error.hpp>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
@@ -168,5 +169,7 @@ inline void knn_merge_parts(const value_t* inK,
   else if (k <= 1024)
     knn_merge_parts_impl<value_idx, value_t, 1024, 8>(
       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
@@ -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