Merge pull request NVIDIA#1376 from allisonvacanti/bug/scan_by_key_modernize/gh.1374

alliepiper · web-flow · commit 1ef03743c5b7 · 2021-02-09T15:33:27.000-05:00
Modernize scan_by_key functors / type deductions.
diff --git a/testing/scan_by_key.cu b/testing/scan_by_key.cu
@@ -1,6 +1,7 @@
 #include <unittest/unittest.h>
 #include <thrust/scan.h>
 #include <thrust/functional.h>
+#include <thrust/iterator/discard_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/iterator/retag.h>
 #include <thrust/random.h>
@@ -540,6 +541,74 @@ void TestScanByKeyMixedTypes(void)
 DECLARE_UNITTEST(TestScanByKeyMixedTypes);
 
 
+template <typename T>
+void TestScanByKeyDiscardOutput(std::size_t n)
+{
+  thrust::host_vector<T> h_keys(n);
+  thrust::default_random_engine rng;
+
+  for (size_t i = 0, k = 0; i < n; i++)
+  {
+    h_keys[i] = static_cast<T>(k);
+    if (rng() % 10 == 0)
+    {
+      k++;
+    }
+  }
+  thrust::device_vector<T> d_keys = h_keys;
+
+  thrust::host_vector<T> h_vals(n);
+  for(size_t i = 0; i < n; i++)
+  {
+    h_vals[i] = static_cast<T>(i % 10);
+  }
+  thrust::device_vector<T> d_vals = h_vals;
+
+  auto out = thrust::make_discard_iterator();
+
+  // These are no-ops, but they should compile.
+  thrust::exclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out);
+  thrust::exclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out,
+                                T{});
+  thrust::exclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out,
+                                T{},
+                                thrust::equal_to<T>{});
+  thrust::exclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out,
+                                T{},
+                                thrust::equal_to<T>{},
+                                thrust::multiplies<T>{});
+
+  thrust::inclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out);
+  thrust::inclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out,
+                                thrust::equal_to<T>{});
+  thrust::inclusive_scan_by_key(d_keys.cbegin(),
+                                d_keys.cend(),
+                                d_vals.cbegin(),
+                                out,
+                                thrust::equal_to<T>{},
+                                thrust::multiplies<T>{});
+}
+DECLARE_VARIABLE_UNITTEST(TestScanByKeyDiscardOutput);
+
+
 void TestScanByKeyLargeInput()
 {
     const unsigned int N = 1 << 20;
diff --git a/thrust/system/cuda/detail/scan_by_key.h b/thrust/system/cuda/detail/scan_by_key.h
@@ -844,14 +844,14 @@ inclusive_scan_by_key(execution_policy<Derived> &policy,
                       ValOutputIt                value_result,
                       BinaryPred                 binary_pred)
 {
-  typedef typename thrust::iterator_traits<ValOutputIt>::value_type value_type;
+  typedef typename thrust::iterator_traits<ValInputIt>::value_type value_type;
   return cuda_cub::inclusive_scan_by_key(policy,
                                          key_first,
                                          key_last,
                                          value_first,
                                          value_result,
                                          binary_pred,
-                                         plus<value_type>());
+                                         thrust::plus<>());
 }
 
 template <class Derived,
@@ -871,7 +871,7 @@ inclusive_scan_by_key(execution_policy<Derived> &policy,
                                          key_last,
                                          value_first,
                                          value_result,
-                                         equal_to<key_type>());
+                                         thrust::equal_to<>());
 }
 
 
@@ -948,7 +948,7 @@ exclusive_scan_by_key(execution_policy<Derived> &policy,
                                          value_result,
                                          init,
                                          binary_pred,
-                                         plus<Init>());
+                                         plus<>());
 }
 
 template <class Derived,
@@ -971,7 +971,7 @@ exclusive_scan_by_key(execution_policy<Derived> &policy,
                                          value_first,
                                          value_result,
                                          init,
-                                         equal_to<key_type>());
+                                         equal_to<>());
 }
 
 
@@ -986,13 +986,13 @@ exclusive_scan_by_key(execution_policy<Derived> &policy,
                       ValInputIt                 value_first,
                       ValOutputIt                value_result)
 {
-  typedef typename iterator_traits<ValOutputIt>::value_type value_type;
+  typedef typename iterator_traits<ValInputIt>::value_type value_type;
   return cuda_cub::exclusive_scan_by_key(policy,
                                          key_first,
                                          key_last,
                                          value_first,
                                          value_result,
-                                         value_type(0));
+                                         value_type{});
 }
 
 
diff --git a/thrust/system/detail/generic/scan_by_key.inl b/thrust/system/detail/generic/scan_by_key.inl
@@ -16,6 +16,7 @@
 
 
 #include <thrust/detail/config.h>
+#include <thrust/detail/cstdint.h>
 #include <thrust/system/detail/generic/scan_by_key.h>
 #include <thrust/functional.h>
 #include <thrust/transform.h>
@@ -71,8 +72,7 @@ __host__ __device__
                                        InputIterator2 first2,
                                        OutputIterator result)
 {
-  typedef typename thrust::iterator_traits<InputIterator1>::value_type InputType1;
-  return thrust::inclusive_scan_by_key(exec, first1, last1, first2, result, thrust::equal_to<InputType1>());
+  return thrust::inclusive_scan_by_key(exec, first1, last1, first2, result, thrust::equal_to<>());
 }
 
 
@@ -108,8 +108,8 @@ __host__ __device__
                                        BinaryPredicate binary_pred,
                                        AssociativeOperator binary_op)
 {
-  typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;
-  typedef unsigned int HeadFlagType;
+  using OutputType = typename thrust::iterator_traits<InputIterator2>::value_type;
+  using HeadFlagType = thrust::detail::uint32_t;
 
   const size_t n = last1 - first1;
 
@@ -146,8 +146,8 @@ __host__ __device__
                                        InputIterator2 first2,
                                        OutputIterator result)
 {
-  typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;
-  return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, OutputType(0));
+  typedef typename thrust::iterator_traits<InputIterator2>::value_type InitType;
+  return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, InitType{});
 }
 
 
@@ -164,8 +164,7 @@ __host__ __device__
                                        OutputIterator result,
                                        T init)
 {
-  typedef typename thrust::iterator_traits<InputIterator1>::value_type InputType1;
-  return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, init, thrust::equal_to<InputType1>());
+  return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, init, thrust::equal_to<>());
 }
 
 
@@ -205,8 +204,8 @@ __host__ __device__
                                        BinaryPredicate binary_pred,
                                        AssociativeOperator binary_op)
 {
-  typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;
-  typedef unsigned int HeadFlagType;
+  using OutputType = T;
+  using HeadFlagType = thrust::detail::uint32_t;
 
   const size_t n = last1 - first1;
 
diff --git a/thrust/system/detail/sequential/scan_by_key.h b/thrust/system/detail/sequential/scan_by_key.h
@@ -52,8 +52,8 @@ __host__ __device__
                                        BinaryPredicate binary_pred,
                                        BinaryFunction binary_op)
 {
-  typedef typename thrust::iterator_traits<InputIterator1>::value_type KeyType;
-  typedef typename thrust::iterator_traits<OutputIterator>::value_type ValueType;
+  using KeyType = typename thrust::iterator_traits<InputIterator1>::value_type;
+  using ValueType = typename thrust::iterator_traits<InputIterator2>::value_type;
 
   // wrap binary_op
   thrust::detail::wrapped_function<
@@ -105,8 +105,8 @@ __host__ __device__
                                        BinaryPredicate binary_pred,
                                        BinaryFunction binary_op)
 {
-  typedef typename thrust::iterator_traits<InputIterator1>::value_type KeyType;
-  typedef typename thrust::iterator_traits<OutputIterator>::value_type ValueType;
+  using KeyType = typename thrust::iterator_traits<InputIterator1>::value_type;
+  using ValueType = T;
 
   if(first1 != last1)
   {

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@`
`16`	`16`
`17`	`17`
`18`	`18`	`#include <thrust/detail/config.h>`
	`19`	`+#include <thrust/detail/cstdint.h>`
`19`	`20`	`#include <thrust/system/detail/generic/scan_by_key.h>`
`20`	`21`	`#include <thrust/functional.h>`
`21`	`22`	`#include <thrust/transform.h>`
`@@ -71,8 +72,7 @@ __host__ __device__`
`71`	`72`	`InputIterator2 first2,`
`72`	`73`	`OutputIterator result)`
`73`	`74`	`{`
`74`		`- typedef typename thrust::iterator_traits<InputIterator1>::value_type InputType1;`
`75`		`- return thrust::inclusive_scan_by_key(exec, first1, last1, first2, result, thrust::equal_to<InputType1>());`
	`75`	`+ return thrust::inclusive_scan_by_key(exec, first1, last1, first2, result, thrust::equal_to<>());`
`76`	`76`	`}`
`77`	`77`
`78`	`78`
`@@ -108,8 +108,8 @@ __host__ __device__`
`108`	`108`	`BinaryPredicate binary_pred,`
`109`	`109`	`AssociativeOperator binary_op)`
`110`	`110`	`{`
`111`		`- typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;`
`112`		`- typedef unsigned int HeadFlagType;`
	`111`	`+ using OutputType = typename thrust::iterator_traits<InputIterator2>::value_type;`
	`112`	`+ using HeadFlagType = thrust::detail::uint32_t;`
`113`	`113`
`114`	`114`	`const size_t n = last1 - first1;`
`115`	`115`
`@@ -146,8 +146,8 @@ __host__ __device__`
`146`	`146`	`InputIterator2 first2,`
`147`	`147`	`OutputIterator result)`
`148`	`148`	`{`
`149`		`- typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;`
`150`		`- return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, OutputType(0));`
	`149`	`+ typedef typename thrust::iterator_traits<InputIterator2>::value_type InitType;`
	`150`	`+ return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, InitType{});`
`151`	`151`	`}`
`152`	`152`
`153`	`153`
`@@ -164,8 +164,7 @@ __host__ __device__`
`164`	`164`	`OutputIterator result,`
`165`	`165`	`T init)`
`166`	`166`	`{`
`167`		`- typedef typename thrust::iterator_traits<InputIterator1>::value_type InputType1;`
`168`		`- return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, init, thrust::equal_to<InputType1>());`
	`167`	`+ return thrust::exclusive_scan_by_key(exec, first1, last1, first2, result, init, thrust::equal_to<>());`
`169`	`168`	`}`
`170`	`169`
`171`	`170`
`@@ -205,8 +204,8 @@ __host__ __device__`
`205`	`204`	`BinaryPredicate binary_pred,`
`206`	`205`	`AssociativeOperator binary_op)`
`207`	`206`	`{`
`208`		`- typedef typename thrust::iterator_traits<OutputIterator>::value_type OutputType;`
`209`		`- typedef unsigned int HeadFlagType;`
	`207`	`+ using OutputType = T;`
	`208`	`+ using HeadFlagType = thrust::detail::uint32_t;`
`210`	`209`
`211`	`210`	`const size_t n = last1 - first1;`
`212`	`211`
Original file line number	Diff line number	Diff line change
`@@ -52,8 +52,8 @@ __host__ __device__`
`52`	`52`	`BinaryPredicate binary_pred,`
`53`	`53`	`BinaryFunction binary_op)`
`54`	`54`	`{`
`55`		`- typedef typename thrust::iterator_traits<InputIterator1>::value_type KeyType;`
`56`		`- typedef typename thrust::iterator_traits<OutputIterator>::value_type ValueType;`
	`55`	`+ using KeyType = typename thrust::iterator_traits<InputIterator1>::value_type;`
	`56`	`+ using ValueType = typename thrust::iterator_traits<InputIterator2>::value_type;`
`57`	`57`
`58`	`58`	`// wrap binary_op`
`59`	`59`	`thrust::detail::wrapped_function<`
`@@ -105,8 +105,8 @@ __host__ __device__`
`105`	`105`	`BinaryPredicate binary_pred,`
`106`	`106`	`BinaryFunction binary_op)`
`107`	`107`	`{`
`108`		`- typedef typename thrust::iterator_traits<InputIterator1>::value_type KeyType;`
`109`		`- typedef typename thrust::iterator_traits<OutputIterator>::value_type ValueType;`
	`108`	`+ using KeyType = typename thrust::iterator_traits<InputIterator1>::value_type;`
	`109`	`+ using ValueType = T;`
`110`	`110`
`111`	`111`	`if(first1 != last1)`
`112`	`112`	`{`