Allow different input value types for merge

As long as they are convertible to the value type of the first iterator. This weakens the publicly documented guarantees of equal value types to restore the old behavior of the thrust implementation replaced in #1817.

Author: bernhardmgruber

cub/cub/agent/agent_merge.cuh

@@ -59,6 +59,7 @@ struct agent_t
 {
   using policy = Policy;
 
+  // key and value type are taken from the first input sequence (consistent with old Thrust behavior)
   using key_type  = typename ::cuda::std::iterator_traits<KeysIt1>::value_type;
   using item_type = typename ::cuda::std::iterator_traits<ItemsIt1>::value_type;
 

cub/cub/agent/agent_merge_sort.cuh

@@ -382,19 +382,20 @@ gmem_to_reg(T (&output)[ITEMS_PER_THREAD], It1 input1, It2 input2, int count1, i
 #pragma unroll
     for (int item = 0; item < ITEMS_PER_THREAD; ++item)
     {
-      int idx      = BLOCK_THREADS * item + threadIdx.x;
-      output[item] = (idx < count1) ? input1[idx] : input2[idx - count1];
+      const int idx = BLOCK_THREADS * item + threadIdx.x;
+      // It1 and It2 could have different value types. Convert after load.
+      output[item] = (idx < count1) ? static_cast<T>(input1[idx]) : static_cast<T>(input2[idx - count1]);
     }
   }
   else
   {
 #pragma unroll
     for (int item = 0; item < ITEMS_PER_THREAD; ++item)
     {
-      int idx = BLOCK_THREADS * item + threadIdx.x;
+      const int idx = BLOCK_THREADS * item + threadIdx.x;
       if (idx < count1 + count2)
       {
-        output[item] = (idx < count1) ? input1[idx] : input2[idx - count1];
+        output[item] = (idx < count1) ? static_cast<T>(input1[idx]) : static_cast<T>(input2[idx - count1]);
       }
     }
   }

cub/cub/device/dispatch/dispatch_merge.cuh

@@ -64,11 +64,6 @@ CUB_DETAIL_KERNEL_ATTRIBUTES void device_partition_merge_path_kernel(
   Offset* merge_partitions,
   CompareOp compare_op)
 {
-  static_assert(
-    ::cuda::std::is_convertible<typename ::cuda::std::__invoke_of<CompareOp, value_t<KeyIt1>, value_t<KeyIt1>>::type,
-                                bool>::value,
-    "Comparison operator must be convertible to bool");
-
   // items_per_tile must be the same of the merge kernel later, so we have to consider whether a fallback agent will be
   // selected for the merge agent that changes the tile size
   constexpr int items_per_tile =
@@ -121,9 +116,12 @@ __launch_bounds__(
     Offset* merge_partitions,
     vsmem_t global_temp_storage)
 {
+  // the merge agent loads keys into a local array of KeyIt1::value_type, on which the comparisons are performed
+  using key_t = value_t<KeyIt1>;
+  static_assert(::cuda::std::__invokable<CompareOp, key_t, key_t>::value,
+                "Comparison operator cannot compare two keys");
   static_assert(
-    ::cuda::std::is_convertible<typename ::cuda::std::__invoke_of<CompareOp, value_t<KeyIt1>, value_t<KeyIt1>>::type,
-                                bool>::value,
+    ::cuda::std::is_convertible<typename ::cuda::std::__invoke_of<CompareOp, key_t, key_t>::type, bool>::value,
     "Comparison operator must be convertible to bool");
 
   using MergeAgent = typename choose_merge_agent<
@@ -218,15 +216,6 @@ template <typename KeyIt1,
           typename PolicyHub = device_merge_policy_hub<value_t<KeyIt1>, value_t<ValueIt1>>>
 struct dispatch_t
 {
-  using key_t   = cub::detail::value_t<KeyIt1>;
-  using value_t = cub::detail::value_t<ValueIt1>;
-
-  // Cannot check output iterators, since they could be discard iterators, which do not have the right value_type
-  static_assert(::cuda::std::is_same<cub::detail::value_t<KeyIt2>, key_t>::value, "");
-  static_assert(::cuda::std::is_same<cub::detail::value_t<ValueIt2>, value_t>::value, "");
-  static_assert(::cuda::std::__invokable<CompareOp, key_t, key_t>::value,
-                "Comparison operator cannot compare two keys");
-
   void* d_temp_storage;
   std::size_t& temp_storage_bytes;
   KeyIt1 d_keys1;
@@ -351,7 +340,7 @@ struct dispatch_t
     {
       return error;
     }
-    dispatch_t dispatch{std::forward<Args>(args)...};
+    dispatch_t dispatch{::cuda::std::forward<Args>(args)...};
     error = CubDebug(PolicyHub::max_policy::Invoke(ptx_version, dispatch));
     if (cudaSuccess != error)
     {