Fix smem index calculation

rapidsai · Mar 25, 2024 · 8c52e25 · 8c52e25
1 parent 43f327c
commit 8c52e25
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/cpp/include/raft/neighbors/detail/cagra/compute_distance_vpq.cuh b/cpp/include/raft/neighbors/detail/cagra/compute_distance_vpq.cuh
@@ -61,10 +61,11 @@ struct cagra_q_dataset_descriptor_t : public dataset_descriptor_base_t<half, DIS
 
         // Change the order of PQ code book array to reduce the
         // frequency of bank conflicts.
-        constexpr auto num_elements_per_bank = 4 / utils::size_of<CODE_BOOK_T>();
-        const auto j                         = i / num_elements_per_bank;
+        constexpr auto num_elements_per_bank  = 4 / utils::size_of<CODE_BOOK_T>();
+        constexpr auto num_banks_per_subspace = PQ_LEN / num_elements_per_bank;
+        const auto j                          = i / num_elements_per_bank;
         const auto smem_index =
-          (j / num_elements_per_bank) + (j % num_elements_per_bank) * (1 << PQ_BITS);
+          (j / num_banks_per_subspace) + (j % num_banks_per_subspace) * (1 << PQ_BITS);
         reinterpret_cast<half2*>(smem_pq_code_book_ptr)[smem_index] = buf2;
       }
     } else {