Fix insert_or_apply bug (#546)

NVIDIA · Jul 16, 2024 · ee5c104 · ee5c104
1 parent e7b5a38
commit ee5c104
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
@@ -605,6 +605,9 @@ class operator_impl<
 
         // If the key is already in the container, update the payload and return
         if (eq_res == detail::equal_result::EQUAL) {
+          if constexpr (sizeof(value_type) > 8) {
+            ref_.impl_.wait_for_payload(slot_ptr->second, empty_value);
+          }
           op(cuda::atomic_ref<T, Scope>{slot_ptr->second}, val.second);
           return;
         }
@@ -689,6 +692,9 @@ class operator_impl<
       if (group_contains_equal) {
         auto const src_lane = __ffs(group_contains_equal) - 1;
         if (group.thread_rank() == src_lane) {
+          if constexpr (sizeof(value_type) > 8) {
+            ref_.impl_.wait_for_payload(slot_ptr->second, empty_value);
+          }
           op(cuda::atomic_ref<T, Scope>{slot_ptr->second}, val.second);
         }
         return;

diff --git a/tests/static_map/insert_or_apply_test.cu b/tests/static_map/insert_or_apply_test.cu
@@ -92,7 +92,7 @@ TEMPLATE_TEST_CASE_SIG(
   (int64_t, int32_t, cuco::test::probe_sequence::linear_probing, 2),
   (int64_t, int64_t, cuco::test::probe_sequence::linear_probing, 2))
 {
-  constexpr size_type num_keys{400};
+  constexpr size_type num_keys{10'000};
   constexpr size_type num_unique_keys{100};
 
   using probe = std::conditional_t<