Use transparent functionals in placeholder expressions.

Fixes and adds regression tests for NVIDIA#1178 & NVIDIA#1229.
alliepiper · Aug 28, 2020 · aa6dc83 · aa6dc83
1 parent ec5baea
commit aa6dc83
Show file tree

Hide file tree

Showing 11 changed files with 604 additions and 501 deletions.
diff --git a/testing/find.cu b/testing/find.cu
@@ -1,4 +1,5 @@
 #include <unittest/unittest.h>
+#include <thrust/sequence.h>
 #include <thrust/find.h>
 #include <thrust/iterator/retag.h>
 
@@ -338,3 +339,35 @@ void TestFindWithBigIndexes()
     TestFindWithBigIndexesHelper(33);
 }
 DECLARE_UNITTEST(TestFindWithBigIndexes);
+
+namespace
+{
+
+class Weird
+{
+  int value;
+
+public:
+  __host__ __device__ Weird(int val, int)
+      : value(val)
+  {}
+
+  friend __host__ __device__
+  bool operator==(int x, Weird y)
+  {
+    return x == y.value;
+  }
+};
+
+} // end anon namespace
+
+void TestFindAsymmetricEquality()
+{ // Regression test for thrust/thrust#1229
+  thrust::host_vector<int> v(1000);
+  thrust::sequence(v.begin(), v.end());
+  thrust::device_vector<int> dv(v);
+  auto result = thrust::find(dv.begin(), dv.end(), Weird(333, 0));
+  ASSERT_EQUAL(*result, 333);
+  ASSERT_EQUAL(result - dv.begin(), 333);
+}
+DECLARE_UNITTEST(TestFindAsymmetricEquality);
diff --git a/testing/inner_product.cu b/testing/inner_product.cu
@@ -1,8 +1,11 @@
 #include <unittest/unittest.h>
 #include <thrust/inner_product.h>
+
+#include <thrust/functional.h>
 #include <thrust/iterator/retag.h>
 #include <thrust/device_malloc.h>
 #include <thrust/device_free.h>
+#include <thrust/device_vector.h>
 
 template <class Vector>
 void TestInnerProductSimple(void)
@@ -153,3 +156,18 @@ void TestInnerProductWithBigIndexes()
     TestInnerProductWithBigIndexesHelper(33);
 }
 DECLARE_UNITTEST(TestInnerProductWithBigIndexes);
+
+void TestInnerProductPlaceholders()
+{ // Regression test for thrust/thrust#1178
+  using namespace thrust::placeholders;
+
+  thrust::device_vector<float> v1(100, 1.f);
+  thrust::device_vector<float> v2(100, 1.f);
+
+  auto result = thrust::inner_product(v1.begin(), v1.end(), v2.begin(), 0.0f,
+                                      thrust::plus<float>{},
+                                      _1 * _2 + 1.0f);
+
+  ASSERT_ALMOST_EQUAL(result, 200.f);
+}
+DECLARE_UNITTEST(TestInnerProductPlaceholders);
diff --git a/thrust/detail/functional/actor.h b/thrust/detail/functional/actor.h
@@ -30,6 +30,7 @@
 #include <thrust/detail/functional/value.h>
 #include <thrust/detail/functional/composite.h>
 #include <thrust/detail/functional/operators/assignment_operator.h>
+#include <thrust/detail/raw_reference_cast.h>
 #include <thrust/detail/type_traits/result_of_adaptable_function.h>
 
 namespace thrust
@@ -39,6 +40,14 @@ namespace detail
 namespace functional
 {
 
+// eval_ref<T> is
+// - T when T is a subclass of thrust::reference
+// - T& otherwise
+// This is used to let thrust::references pass through actor evaluations.
+template <typename T>
+using eval_ref = typename std::conditional<
+  thrust::detail::is_wrapped_reference<T>::value, T, T&>::type;
+
 template<typename Action, typename Env>
   struct apply_actor
 {
@@ -61,55 +70,10 @@ template<typename Eval>
   typename apply_actor<eval_type, thrust::null_type >::type
   operator()(void) const;
 
-  template<typename T0>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&> >::type
-  operator()(T0 &_0) const;
-
-  template<typename T0, typename T1>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&> >::type
-  operator()(T0 &_0, T1 &_1) const;
-
-  template<typename T0, typename T1, typename T2>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2) const;
-
-  template<typename T0, typename T1, typename T2, typename T3>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8>
-  __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&,T8&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7, T8 &_8) const;
-
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
+  template <typename... Ts>
   __host__ __device__
-  typename apply_actor<eval_type, thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&,T8&,T9&> >::type
-  operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7, T8 &_8, T9 &_9) const;
+  typename apply_actor<eval_type, thrust::tuple<eval_ref<Ts>...>>::type
+  operator()(Ts&&... ts) const;
 
   template<typename T>
   __host__ __device__

diff --git a/thrust/detail/functional/actor.inl b/thrust/detail/functional/actor.inl
@@ -27,6 +27,9 @@
 #include <thrust/detail/functional/composite.h>
 #include <thrust/detail/functional/operators/assignment_operator.h>
 #include <thrust/functional.h>
+#include <thrust/type_traits/logical_metafunctions.h>
+
+#include <type_traits>
 
 namespace thrust
 {
@@ -62,135 +65,38 @@ template<typename Eval>
   return eval_type::eval(thrust::null_type());
 } // end basic_environment::operator()
 
-template<typename Eval>
-  template<typename T0>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0) const
-{
-  return eval_type::eval(thrust::tie(_0));
-} // end basic_environment::operator()
+// actor::operator() needs to construct a tuple of references to its
+// arguments. To make this work with thrust::reference<T>, we need to
+// detect thrust proxy references and store them as T rather than T&.
+// This check ensures that the forwarding references passed into
+// actor::operator() are either:
+// - T&& if and only if T is a thrust::reference<U>, or
+// - T& for any other types.
+// This struct provides a nicer diagnostic for when these conditions aren't
+// met.
+template <typename T>
+using actor_check_ref_type =
+  thrust::detail::integral_constant<bool,
+    ( std::is_lvalue_reference<T>::value ||
+      thrust::detail::is_wrapped_reference<T>::value )>;
+
+template <typename... Ts>
+using actor_check_ref_types =
+  thrust::conjunction<actor_check_ref_type<Ts>...>;
 
 template<typename Eval>
-  template<typename T0, typename T1>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1) const
+template<typename... Ts>
+__host__ __device__
+typename apply_actor<typename actor<Eval>::eval_type,
+                     thrust::tuple<eval_ref<Ts>...>>::type
+actor<Eval>::operator()(Ts&&... ts) const
 {
-  return eval_type::eval(thrust::tie(_0,_1));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4,_5));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4,_5,_6));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4,_5,_6,_7));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&,T8&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7, T8 &_8) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4,_5,_6,_7,_8));
-} // end basic_environment::operator()
-
-template<typename Eval>
-  template<typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
-    __host__ __device__
-    typename apply_actor<
-      typename actor<Eval>::eval_type,
-      typename thrust::tuple<T0&,T1&,T2&,T3&,T4&,T5&,T6&,T7&,T8&,T9&>
-    >::type
-      actor<Eval>
-        ::operator()(T0 &_0, T1 &_1, T2 &_2, T3 &_3, T4 &_4, T5 &_5, T6 &_6, T7 &_7, T8 &_8, T9 &_9) const
-{
-  return eval_type::eval(thrust::tie(_0,_1,_2,_3,_4,_5,_6,_7,_8,_9));
-} // end basic_environment::operator()
+  static_assert(actor_check_ref_types<Ts...>::value,
+                "Actor evaluations only support rvalue references to "
+                "thrust::reference subclasses.");
+  using tuple_type = thrust::tuple<eval_ref<Ts>...>;
+  return eval_type::eval(tuple_type(THRUST_FWD(ts)...));
+} // end actor<Eval>::operator()
 
 template<typename Eval>
   template<typename T>