Add zip_function to adapt N-ary functions to take a tuple

bjude · alliepiper · commit 82370c21670b · 2020-03-04T14:37:49.000-05:00
Eases the use general function objects with zip iterators without modifying them or hand writing a wrapping class

Test for zip_function

Based on the zip iterator transform test

zip_function: Move details into thrust::detal::zip_detail

zip_function: make operator() const and make stored function mutable

CMake: Add filter for test that require c++11

Only add zip_function for now, making the list exhaustive can be another PR

zip_function: Add example to arbitrary_transformation

zip_function: Add c++11 guard

zip_function: Documentation

Zip Function: newline at end of file

Allison rewrote some bits to support C++11 compilers.

Reviewed-by: Allison Vacanti &lt;alliepiper16@gmail.com&gt;
diff --git a/examples/arbitrary_transformation.cu b/examples/arbitrary_transformation.cu
@@ -3,6 +3,12 @@
 #include <thrust/iterator/zip_iterator.h>
 #include <iostream>
 
+#include <thrust/detail/config.h>
+
+#if THRUST_CPP_DIALECT >= 2011 && !defined(THRUST_LEGACY_GCC)
+#include <thrust/zip_function.h>
+#endif // >= C++11
+
 // This example shows how to implement an arbitrary transformation of
 // the form output[i] = F(first[i], second[i], third[i], ... ).
 // In this example, we use a function with 3 inputs and 1 output.
@@ -22,6 +28,10 @@
 //      D[i] = A[i] + B[i] * C[i];
 // by invoking arbitrary_functor() on each of the tuples using for_each.
 //
+// If we are using a functor that is not designed for zip iterators by taking a
+// tuple instead of individual arguments we can adapt this function using the
+// zip_function adaptor (C++11 only).
+//
 // Note that we could extend this example to implement functions with an
 // arbitrary number of input arguments by zipping more sequence together.
 // With the same approach we can have multiple *output* sequences, if we 
@@ -31,7 +41,7 @@
 //
 // The possibilities are endless! :)
 
-struct arbitrary_functor
+struct arbitrary_functor1
 {
     template <typename Tuple>
     __host__ __device__
@@ -42,14 +52,25 @@ struct arbitrary_functor
     }
 };
 
+#if THRUST_CPP_DIALECT >= 2011 && !defined(THRUST_LEGACY_GCC)
+struct arbitrary_functor2
+{
+    __host__ __device__
+    void operator()(const float& a, const float& b, const float& c, float& d)
+    {
+        // D[i] = A[i] + B[i] * C[i];
+        d = a + b * c;
+    }
+};
+#endif // >= C++11
 
 int main(void)
 {
     // allocate storage
     thrust::device_vector<float> A(5);
     thrust::device_vector<float> B(5);
     thrust::device_vector<float> C(5);
-    thrust::device_vector<float> D(5);
+    thrust::device_vector<float> D1(5);
 
     // initialize input vectors
     A[0] = 3;  B[0] = 6;  C[0] = 2; 
@@ -59,12 +80,26 @@ int main(void)
     A[4] = 2;  B[4] = 8;  C[4] = 3; 
 
     // apply the transformation
-    thrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin(), D.begin())),
-                     thrust::make_zip_iterator(thrust::make_tuple(A.end(),   B.end(),   C.end(),   D.end())),
-                     arbitrary_functor());
+    thrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin(), D1.begin())),
+                     thrust::make_zip_iterator(thrust::make_tuple(A.end(),   B.end(),   C.end(),   D1.end())),
+                     arbitrary_functor1());
+
+    // print the output
+    std::cout << "Tuple functor" << std::endl;
+    for(int i = 0; i < 5; i++)
+        std::cout << A[i] << " + " << B[i] << " * " << C[i] << " = " << D1[i] << std::endl;
+
+    // apply the transformation using zip_function
+#if THRUST_CPP_DIALECT >= 2011 && !defined(THRUST_LEGACY_GCC)
+    thrust::device_vector<float> D2(5);
+    thrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin(), D2.begin())),
+                     thrust::make_zip_iterator(thrust::make_tuple(A.end(),   B.end(),   C.end(),   D2.end())),
+                     thrust::make_zip_function(arbitrary_functor2()));
 
     // print the output
+    std::cout << "N-ary functor" << std::endl;
     for(int i = 0; i < 5; i++)
-        std::cout << A[i] << " + " << B[i] << " * " << C[i] << " = " << D[i] << std::endl;
+        std::cout << A[i] << " + " << B[i] << " * " << C[i] << " = " << D2[i] << std::endl;
+#endif // >= C++11
 }
 
diff --git a/testing/zip_function.cu b/testing/zip_function.cu
@@ -0,0 +1,70 @@
+#include <thrust/detail/config.h>
+
+#if THRUST_CPP_DIALECT >= 2011 && !defined(THRUST_LEGACY_GCC)
+
+#include <unittest/unittest.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/transform.h>
+#include <thrust/zip_function.h>
+
+#include <iostream>
+
+using namespace unittest;
+
+struct SumThree
+{
+  template <typename T1, typename T2, typename T3>
+  __host__ __device__
+  auto operator()(T1 x, T2 y, T3 z) const
+  THRUST_DECLTYPE_RETURNS(x + y + z)
+}; // end SumThree
+
+struct SumThreeTuple
+{
+  template <typename Tuple>
+  __host__ __device__
+  auto operator()(Tuple x) const
+  THRUST_DECLTYPE_RETURNS(thrust::get<0>(x) + thrust::get<1>(x) + thrust::get<2>(x))
+}; // end SumThreeTuple
+
+template <typename T>
+struct TestZipFunctionTransform
+{
+  void operator()(const size_t n)
+  {
+    using namespace thrust;
+
+    host_vector<T> h_data0 = unittest::random_samples<T>(n);
+    host_vector<T> h_data1 = unittest::random_samples<T>(n);
+    host_vector<T> h_data2 = unittest::random_samples<T>(n);
+
+    device_vector<T> d_data0 = h_data0;
+    device_vector<T> d_data1 = h_data1;
+    device_vector<T> d_data2 = h_data2;
+
+    host_vector<T>   h_result_tuple(n);
+    host_vector<T>   h_result_zip(n);
+    device_vector<T> d_result_zip(n);
+
+    // Tuple base case
+    transform(make_zip_iterator(make_tuple(h_data0.begin(), h_data1.begin(), h_data2.begin())),
+              make_zip_iterator(make_tuple(h_data0.end(),   h_data1.end(),   h_data2.end())),
+              h_result_tuple.begin(),
+              SumThreeTuple{});
+    // Zip Function
+    transform(make_zip_iterator(make_tuple(h_data0.begin(), h_data1.begin(), h_data2.begin())),
+              make_zip_iterator(make_tuple(h_data0.end(),   h_data1.end(),   h_data2.end())),
+              h_result_zip.begin(),
+              make_zip_function(SumThree{}));
+    transform(make_zip_iterator(make_tuple(d_data0.begin(), d_data1.begin(), d_data2.begin())),
+              make_zip_iterator(make_tuple(d_data0.end(),   d_data1.end(),   d_data2.end())),
+              d_result_zip.begin(),
+              make_zip_function(SumThree{}));
+
+    ASSERT_EQUAL(h_result_tuple, h_result_zip);
+    ASSERT_EQUAL(h_result_tuple, d_result_zip);
+  }
+};
+VariableUnitTest<TestZipFunctionTransform, ThirtyTwoBitTypes> TestZipFunctionTransformInstance;
+
+#endif // THRUST_CPP_DIALECT
diff --git a/thrust/zip_function.h b/thrust/zip_function.h
@@ -0,0 +1,209 @@
+
+/*! \file thrust/zip_function.h
+ *  \brief Adaptor type that turns an N-ary function object into one that takes
+ *         a tuple of size N so it can easily be used with algorithms taking zip
+ *         iterators
+ */
+
+#pragma once
+
+#include <thrust/detail/config.h>
+#include <thrust/detail/cpp11_required.h>
+#include <thrust/detail/modern_gcc_required.h>
+
+#if THRUST_CPP_DIALECT >= 2011 && !defined(THRUST_LEGACY_GCC)
+
+#include <thrust/type_traits/integer_sequence.h>
+#include <thrust/detail/type_deduction.h>
+
+THRUST_BEGIN_NS
+
+/*! \addtogroup function_objects Function Objects
+ *  \{
+ */
+
+/*! \addtogroup function_object_adaptors Function Object Adaptors
+ *  \ingroup function_objects
+ *  \{
+ */
+
+namespace detail {
+namespace zip_detail {
+
+// Add workaround for decltype(auto) on C++11-only compilers:
+#if THRUST_CPP_DIALECT >= 2014
+
+template <typename Function, typename Tuple, std::size_t... Is>
+__host__ __device__
+decltype(auto) apply_impl(Function&& func, Tuple&& args, index_sequence<Is...>)
+{
+  return func(thrust::get<Is>(THRUST_FWD(args))...);
+}
+
+template <typename Function, typename Tuple>
+__host__ __device__
+decltype(auto) apply(Function&& func, Tuple&& args)
+{
+  constexpr auto tuple_size = thrust::tuple_size<typename std::decay<Tuple>::type>::value;
+  return apply_impl(THRUST_FWD(func), THRUST_FWD(args), make_index_sequence<tuple_size>{});
+}
+
+#else // THRUST_CPP_DIALECT
+
+template <typename Function, typename Tuple, std::size_t... Is>
+__host__ __device__
+auto apply_impl(Function&& func, Tuple&& args, index_sequence<Is...>)
+THRUST_DECLTYPE_RETURNS(func(thrust::get<Is>(THRUST_FWD(args))...))
+
+template <typename Function, typename Tuple>
+__host__ __device__
+auto apply(Function&& func, Tuple&& args)
+THRUST_DECLTYPE_RETURNS(
+    apply_impl(
+      THRUST_FWD(func),
+      THRUST_FWD(args),
+      make_index_sequence<
+        thrust::tuple_size<typename std::decay<Tuple>::type>::value>{})
+)
+
+#endif // THRUST_CPP_DIALECT
+
+} // namespace zip_detail
+} // namespace detail
+
+/*! \p zip_function is a function object that allows the easy use of N-ary 
+ *  function objects with \p zip_iterators without redefining them to take a
+ *  \p tuple instead of N arguments.
+ *
+ *  This means that if a functor that takes 2 arguments which could be used with
+ *  the \p transform function and \p device_iterators can be extended to take 3
+ *  arguments and \p zip_iterators without rewriting the functor in terms of
+ *  \p tuple.
+ * 
+ *  The \p make_zip_function convenience function is provided to avoid having
+ *  to explicitely define the type of the functor when creating a \p zip_function, 
+ *  whic is especially helpful when using lambdas as the functor.
+ *  
+ *  \code
+ *  #include <thrust/iterator/zip_iterator.h>
+ *  #include <thrust/device_vector.h>
+ *  #include <thrust/transform.h>
+ *  #include <thrust/zip_function.h>
+ * 
+ *  struct SumTuple {
+ *    float operator()(Tuple tup) {
+ *      return std::get<0>(tup) + std::get<1>(tup) + std::get<2>(tup);
+ *    }
+ *  };
+ *  struct SumArgs {
+ *    float operator()(float a, float b, float c) {
+ *      return a + b + c;
+ *    }
+ *  };
+ *  
+ *  int main() {
+ *    thrust::device_vector<float> A(3);
+ *    thrust::device_vector<float> B(3);
+ *    thrust::device_vector<float> C(3);
+ *    thrust::device_vector<float> D(3);
+ *    A[0] = 0.f; A[1] = 1.f; A[2] = 2.f;
+ *    B[0] = 1.f; B[1] = 2.f; B[2] = 3.f;
+ *    C[0] = 2.f; C[1] = 3.f; C[2] = 4.f;
+ * 
+ *    // The following four invocations of transform are equivalent
+ *    // Transform with 3-tuple
+ *    thrust::transform(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin())),
+ *                      thrust::make_zip_iterator(thrust::make_tuple(A.end(), B.end(), C.end())),
+ *                      D.begin(),
+ *                      SumTuple{});
+ * 
+ *    // Transform with 3 parameters
+ *    thrust::zip_function<SumArgs> adapted{};
+ *    thrust::transform(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin())),
+ *                      thrust::make_zip_iterator(thrust::make_tuple(A.end(), B.end(), C.end())),
+ *                      D.begin(),
+ *                      adapted);
+ * 
+ *    // Transform with 3 parameters with convenience function
+ *    thrust::zip_function<SumArgs> adapted{};
+ *    thrust::transform(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin())),
+ *                      thrust::make_zip_iterator(thrust::make_tuple(A.end(), B.end(), C.end())),
+ *                      D.begin(),
+ *                      thrust::make_zip_function(SumArgs{}));
+ * 
+ *    // Transform with 3 parameters with convenience function and lambda
+ *    thrust::zip_function<SumArgs> adapted{};
+ *    thrust::transform(thrust::make_zip_iterator(thrust::make_tuple(A.begin(), B.begin(), C.begin())),
+ *                      thrust::make_zip_iterator(thrust::make_tuple(A.end(), B.end(), C.end())),
+ *                      D.begin(),
+ *                      thrust::make_zip_function([] (float a, float b, float c) {
+ *                                                  return a + b + c;
+ *                                                }));
+ *    return 0;
+ *  }
+ *  \endcode
+ * 
+ *  \see make_zip_function
+ *  \see zip_iterator
+ */
+template <typename Function>
+class zip_function
+{
+  public:
+     __host__ __device__
+    zip_function(Function func) : func(std::move(func)) {}
+
+// Add workaround for decltype(auto) on C++11-only compilers:
+#if THRUST_CPP_DIALECT >= 2014
+
+    template <typename Tuple>
+    __host__ __device__
+    decltype(auto) operator()(Tuple&& args) const
+    {
+        return detail::zip_detail::apply(func, THRUST_FWD(args));
+    }
+
+#else // THRUST_CPP_DIALECT
+
+    // Can't just use THRUST_DECLTYPE_RETURNS here since we need to use
+    // std::declval for the signature components:
+    template <typename Tuple>
+    __host__ __device__
+    auto operator()(Tuple&& args) const
+    noexcept(noexcept(detail::zip_detail::apply(std::declval<Function>(), THRUST_FWD(args))))
+    -> decltype(detail::zip_detail::apply(std::declval<Function>(), THRUST_FWD(args)))
+
+    {
+        return detail::zip_detail::apply(func, THRUST_FWD(args));
+    }
+
+#endif // THRUST_CPP_DIALECT
+
+  private:
+    mutable Function func;
+}; 
+
+/*! \p make_zip_function creates a \p zip_function from a function object.
+ *
+ *  \param fun The N-ary function object.
+ *  \return A \p zip_function that takes a N-tuple.
+ *
+ *  \see zip_function
+ */
+template <typename Function>
+__host__ __device__
+auto make_zip_function(Function&& fun) -> zip_function<typename std::decay<Function>::type>
+{
+    using func_t = typename std::decay<Function>::type;
+    return zip_function<func_t>(THRUST_FWD(fun));
+}
+
+/*! \} // end function_object_adaptors
+ */
+
+/*! \} // end function_objects
+ */
+
+THRUST_END_NS
+
+#endif
