CAGRA-Q search

cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp

@@ -205,6 +205,21 @@ struct standard_dataset_descriptor_t
   static const std::uint32_t smem_buffer_size_in_byte = 0;
   __device__ void set_smem_ptr(void* const){};
 
+  template <uint32_t DATASET_BLOCK_DIM>
+  __device__ void copy_query(const DATA_T* const dmem_query_ptr,
+                             QUERY_T* const smem_query_ptr,
+                             const std::uint32_t query_smem_buffer_length)
+  {
+    for (unsigned i = threadIdx.x; i < query_smem_buffer_length; i += blockDim.x) {
+      unsigned j = device::swizzling(i);
+      if (i < dim) {
+        smem_query_ptr[j] = spatial::knn::detail::utils::mapping<float>{}(dmem_query_ptr[i]);
+      } else {
+        smem_query_ptr[j] = 0.0;
+      }
+    }
+  }
+
   template <uint32_t DATASET_BLOCK_DIM, uint32_t TEAM_SIZE>
   __device__ DISTANCE_T compute_similarity(const QUERY_T* const query_ptr,
                                            const INDEX_T dataset_i,

cpp/include/raft/neighbors/detail/cagra/compute_distance_vpq.cuh

@@ -88,6 +88,26 @@ struct cagra_q_dataset_descriptor_t : public dataset_descriptor_base_t<half, DIS
   {
   }
 
+  template <uint32_t DATASET_BLOCK_DIM>
+  __device__ void copy_query(const DATA_T* const dmem_query_ptr,
+                             QUERY_T* const smem_query_ptr,
+                             const std::uint32_t query_smem_buffer_length)
+  {
+    for (unsigned i = threadIdx.x * 2; i < dim; i += blockDim.x * 2) {
+      half2 buf2 = {CUDART_ZERO_FP16, CUDART_ZERO_FP16};
+      if (i < dim) { buf2.x = static_cast<half>(static_cast<float>(dmem_query_ptr[i])); }
+      if (i + 1 < dim) { buf2.y = static_cast<half>(static_cast<float>(dmem_query_ptr[i + 1])); }
+      if ((PQ_BITS == 8) && (PQ_LEN % 2 == 0)) {
+        // Use swizzling in the condition to reduce bank conflicts in shared
+        // memory, which are likely to occur when pq_code_book_dim is large.
+        ((half2*)smem_query_ptr)[device::swizzling<std::uint32_t, DATASET_BLOCK_DIM / 2>(i / 2)] =
+          buf2;
+      } else {
+        (reinterpret_cast<half2*>(smem_query_ptr + i))[0] = buf2;
+      }
+    }
+  }
+
   template <uint32_t DATASET_BLOCK_DIM, uint32_t TEAM_SIZE>
   __device__ DISTANCE_T compute_similarity(const QUERY_T* const query_ptr,
                                            const INDEX_T node_id,

cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh

@@ -213,14 +213,9 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel(
     }
 #endif
   const DATA_T* const query_ptr = queries_ptr + (dataset_desc.dim * query_id);
-  for (unsigned i = threadIdx.x; i < query_smem_buffer_length; i += blockDim.x) {
-    unsigned j = device::swizzling(i);
-    if (i < dataset_desc.dim) {
-      query_buffer[j] = spatial::knn::detail::utils::mapping<float>{}(query_ptr[i]);
-    } else {
-      query_buffer[j] = 0.0;
-    }
-  }
+  dataset_desc.template copy_query<DATASET_BLOCK_DIM>(
+    query_ptr, query_buffer, query_smem_buffer_length);
+
   if (threadIdx.x == 0) { terminate_flag[0] = 0; }
   INDEX_T* const local_visited_hashmap_ptr =
     visited_hashmap_ptr + (hashmap::get_size(hash_bitlen) * query_id);

cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh

@@ -549,14 +549,9 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel(
   auto filter_flag = terminate_flag;
 
   const DATA_T* const query_ptr = queries_ptr + query_id * dataset_desc.dim;
-  for (unsigned i = threadIdx.x; i < query_smem_buffer_length; i += blockDim.x) {
-    unsigned j = device::swizzling(i);
-    if (i < dataset_desc.dim) {
-      query_buffer[j] = spatial::knn::detail::utils::mapping<float>{}(query_ptr[i]);
-    } else {
-      query_buffer[j] = 0.0;
-    }
-  }
+  dataset_desc.template copy_query<DATASET_BLOCK_DIM>(
+    query_ptr, query_buffer, query_smem_buffer_length);
+
   if (threadIdx.x == 0) {
     terminate_flag[0] = 0;
     topk_ws[0]        = ~0u;