Fix problems with heterogeneous insert input types

include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh

@@ -401,13 +401,7 @@ class open_addressing_ref_impl {
         if (eq_res == detail::equal_result::EMPTY or
             cuco::detail::bitwise_compare(this->extract_key(window_slots[i]),
                                           this->erased_key_sentinel())) {
-          switch ([&]() {
-            if constexpr (sizeof(value_type) <= 8) {
-              return packed_cas(window_ptr + i, window_slots[i], value);
-            } else {
-              return cas_dependent_write(window_ptr + i, window_slots[i], value);
-            }
-          }()) {
+          switch (this->attempt_insert_stable(window_ptr + i, window_slots[i], value)) {
             case insert_result::SUCCESS: {
               return {iterator{&window_ptr[i]}, true};
             }
@@ -485,11 +479,7 @@ class open_addressing_ref_impl {
         auto const res      = group.shfl(reinterpret_cast<intptr_t>(slot_ptr), src_lane);
         auto const status   = [&, target_idx = intra_window_index]() {
           if (group.thread_rank() != src_lane) { return insert_result::CONTINUE; }
-          if constexpr (sizeof(value_type) <= 8) {
-            return packed_cas(slot_ptr, window_slots[target_idx], value);
-          } else {
-            return cas_dependent_write(slot_ptr, window_slots[target_idx], value);
-          }
+          return this->attempt_insert_stable(slot_ptr, window_slots[target_idx], value);
         }();
 
         switch (group.shfl(status, src_lane)) {
@@ -1054,6 +1044,35 @@ class open_addressing_ref_impl {
     }
   }
 
+  /**
+   * @brief Attempts to insert an element into a slot.
+   *
+   * @note Dispatches the correct implementation depending on the container
+   * type and presence of other operator mixins.
+   *
+   * @note `stable` here means that the payload will only be updated once from the sentinel value to
+   * the payload value
+   *
+   * @tparam Value Input type which is implicitly convertible to 'value_type'
+   *
+   * @param address Pointer to the slot in memory
+   * @param expected Element to compare against
+   * @param desired Element to insert
+   *
+   * @return Result of this operation, i.e., success/continue/duplicate
+   */
+  template <typename Value>
+  [[nodiscard]] __device__ insert_result attempt_insert_stable(value_type* address,
+                                                               value_type const& expected,
+                                                               Value const& desired) noexcept
+  {
+    if constexpr (sizeof(value_type) <= 8) {
+      return packed_cas(address, expected, desired);
+    } else {
+      return cas_dependent_write(address, expected, desired);
+    }
+  }
+
   // TODO: Clean up the sentinel handling since it's duplicated in ref and equal wrapper
   value_type empty_slot_sentinel_;  ///< Sentinel value indicating an empty slot
   detail::equal_wrapper<key_type, key_equal> predicate_;  ///< Key equality binary callable

include/cuco/detail/static_map/static_map_ref.inl

@@ -18,6 +18,8 @@
 
 #include <cuco/operator.hpp>
 
+#include <thrust/tuple.h>
+
 #include <cuda/atomic>
 
 #include <cooperative_groups.h>
@@ -248,7 +250,7 @@ class operator_impl<
     static_assert(cg_size == 1, "Non-CG operation is incompatible with the current probing scheme");
 
     ref_type& ref_       = static_cast<ref_type&>(*this);
-    auto const key       = value.first;
+    auto const key       = thrust::get<0>(thrust::raw_reference_cast(value));
     auto& probing_scheme = ref_.impl_.probing_scheme();
     auto storage_ref     = ref_.impl_.storage_ref();
     auto probing_iter    = probing_scheme(key, storage_ref.window_extent());
@@ -264,7 +266,7 @@ class operator_impl<
           auto const intra_window_index = thrust::distance(window_slots.begin(), &slot_content);
           ref_.impl_.atomic_store(
             &((storage_ref.data() + *probing_iter)->data() + intra_window_index)->second,
-            value.second);
+            static_cast<T>(thrust::get<1>(value)));
           return;
         }
         if (eq_res == detail::equal_result::EMPTY or
@@ -297,7 +299,7 @@ class operator_impl<
   {
     ref_type& ref_ = static_cast<ref_type&>(*this);
 
-    auto const key       = value.first;
+    auto const key       = thrust::get<0>(thrust::raw_reference_cast(value));
     auto& probing_scheme = ref_.impl_.probing_scheme();
     auto storage_ref     = ref_.impl_.storage_ref();
     auto probing_iter    = probing_scheme(group, key, storage_ref.window_extent());
@@ -332,7 +334,7 @@ class operator_impl<
         if (group.thread_rank() == src_lane) {
           ref_.impl_.atomic_store(
             &((storage_ref.data() + *probing_iter)->data() + intra_window_index)->second,
-            value.second);
+            static_cast<T>(thrust::get<1>(value)));
         }
         group.sync();
         return;
@@ -377,15 +379,17 @@ class operator_impl<
     ref_type& ref_          = static_cast<ref_type&>(*this);
     auto const expected_key = ref_.impl_.empty_slot_sentinel().first;
 
-    auto old_key      = ref_.impl_.compare_and_swap(&slot->first, expected_key, value.first);
+    auto old_key = ref_.impl_.compare_and_swap(
+      &slot->first, expected_key, static_cast<key_type>(thrust::get<0>(value)));
     auto* old_key_ptr = reinterpret_cast<key_type*>(&old_key);
 
     // if key success or key was already present in the map
     if (cuco::detail::bitwise_compare(*old_key_ptr, expected_key) or
-        (ref_.impl_.predicate().equal_to(*old_key_ptr, value.first) ==
+        (ref_.impl_.predicate().equal_to(*old_key_ptr,
+                                         thrust::get<0>(thrust::raw_reference_cast(value))) ==
          detail::equal_result::EQUAL)) {
       // Update payload
-      ref_.impl_.atomic_store(&slot->second, value.second);
+      ref_.impl_.atomic_store(&slot->second, static_cast<T>(thrust::get<1>(value)));
       return true;
     }
     return false;