Add cub::DeviceTransform N->M API entrypoint (#7473)

Author: bernhardmgruber

cub/cub/device/device_transform.cuh

@@ -108,7 +108,110 @@ private:
     }
   }
 
+  // TODO(bgruber): we want to eventually forward the output tuple to the kernel and optimize writing multiple streams
+  template <detail::transform::requires_stable_address StableAddress = detail::transform::requires_stable_address::no,
+            typename... RandomAccessIteratorsIn,
+            typename... RandomAccessIteratorsOut,
+            typename NumItemsT,
+            typename Predicate,
+            typename TransformOp,
+            typename Env>
+  CUB_RUNTIME_FUNCTION static cudaError_t TransformInternal(
+    ::cuda::std::tuple<RandomAccessIteratorsIn...> inputs,
+    ::cuda::std::tuple<RandomAccessIteratorsOut...> outputs,
+    NumItemsT num_items,
+    Predicate predicate,
+    TransformOp transform_op,
+    Env env)
+  {
+    return TransformInternal<StableAddress>(
+      ::cuda::std::move(inputs),
+      ::cuda::make_zip_iterator(::cuda::std::move(outputs)),
+      num_items,
+      ::cuda::std::move(predicate),
+      ::cuda::std::move(transform_op),
+      ::cuda::std::move(env));
+  }
+
 public:
+  //! @rst
+  //! Overview
+  //! +++++++++++++++++++++++++++++++++++++++++++++
+  //! Transforms many input sequences into many output sequence, by applying a transformation operation on corresponding
+  //! input elements and writing the tuple result to the corresponding output elements. No guarantee is given on the
+  //! identity (i.e. address) of the objects passed to the call operator of the transformation operation.
+  //!
+  //! A Simple Example
+  //! +++++++++++++++++++++++++++++++++++++++++++++
+  //!
+  //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu
+  //!     :language: c++
+  //!     :dedent:
+  //!     :start-after: example-begin transform-many-many
+  //!     :end-before: example-end transform-many-many
+  //!
+  //! @endrst
+  //!
+  //! @param inputs A tuple of iterators to the input sequences where num_items elements are read from each. The
+  //! iterators' value types must be trivially relocatable.
+  //! @param outputs A tuple of iterators to the output sequences where num_items results are written to each. Each
+  //! sequence may point to the beginning of one of the input sequences, performing the transformation inplace. Any
+  //! output sequence must not overlap with any of the input sequence in any other way.
+  //! @param num_items The number of elements in each input and output sequence.
+  //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value
+  //! types must be convertible to the parameters of the function object's call operator. The return type of the call
+  //! operator must be a tuple where each tuple element is assignable to the corresponding dereferenced output
+  //! iterators.
+  //! @param env Execution environment, or cudaStream_t. Default is ``cuda::std::execution::env{}``, which will run on
+  //! stream\ :sub:`0`
+  template <typename... RandomAccessIteratorsIn,
+            typename... RandomAccessIteratorsOut,
+            typename NumItemsT,
+            typename TransformOp,
+            typename Env = ::cuda::std::execution::env<>>
+  CUB_RUNTIME_FUNCTION static cudaError_t Transform(
+    ::cuda::std::tuple<RandomAccessIteratorsIn...> inputs,
+    ::cuda::std::tuple<RandomAccessIteratorsOut...> outputs,
+    NumItemsT num_items,
+    TransformOp transform_op,
+    Env env = {})
+  {
+    _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::Transform");
+    return TransformInternal(
+      ::cuda::std::move(inputs),
+      ::cuda::std::move(outputs),
+      num_items,
+      detail::transform::always_true_predicate{},
+      ::cuda::std::move(transform_op),
+      ::cuda::std::move(env));
+  }
+
+#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
+  // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs.
+  template <typename... RandomAccessIteratorsIn,
+            typename... RandomAccessIteratorsOut,
+            typename NumItemsT,
+            typename TransformOp>
+  CUB_RUNTIME_FUNCTION static cudaError_t Transform(
+    void* d_temp_storage,
+    size_t& temp_storage_bytes,
+    ::cuda::std::tuple<RandomAccessIteratorsIn...> inputs,
+    ::cuda::std::tuple<RandomAccessIteratorsOut...> outputs,
+    NumItemsT num_items,
+    TransformOp transform_op,
+    cudaStream_t stream = nullptr)
+  {
+    if (d_temp_storage == nullptr)
+    {
+      temp_storage_bytes = 1;
+      return cudaSuccess;
+    }
+
+    return Transform(
+      ::cuda::std::move(inputs), ::cuda::std::move(outputs), num_items, ::cuda::std::move(transform_op), stream);
+  }
+#endif // _CCCL_DOXYGEN_INVOKED
+
   //! @rst
   //! Overview
   //! +++++++++++++++++++++++++++++++++++++++++++++

cub/test/catch2_test_device_transform.cu

@@ -435,22 +435,42 @@ C2H_TEST("DeviceTransform::Transform fancy output iterator type with void value
   REQUIRE(result == c2h::device_vector<type>(num_items, 3));
 }
 
-C2H_TEST("DeviceTransform::Transform mixed input iterator types", "[device][transform]")
+struct plus_mul_neg
 {
-  using type          = int;
+  template <typename T>
+  __host__ __device__ auto operator()(T a, T b) const
+  {
+    return cuda::std::tuple{a + b, a * b, -a};
+  }
+};
+
+C2H_TEST("DeviceTransform::Transform mixed iterator types 2 -> 3", "[device][transform]")
+{
+  using type          = unsigned; // overflow is defined
   const int num_items = GENERATE(100, 100'000); // try to hit the small and full tile code paths
   cuda::counting_iterator<type> a{0};
   c2h::device_vector<type> b(num_items, thrust::no_init);
   c2h::gen(C2H_SEED(1), b);
 
-  c2h::device_vector<type> result(num_items, thrust::no_init);
-  transform_many(cuda::std::make_tuple(a, b.begin()), result.begin(), num_items, cuda::std::plus<type>{});
+  c2h::device_vector<type> result_a(num_items, thrust::no_init);
+  c2h::device_vector<type> result_b(num_items, thrust::no_init);
+  c2h::device_vector<type> result_c(num_items, thrust::no_init);
+  transform_many(
+    cuda::std::make_tuple(a, b.begin()),
+    cuda::std::make_tuple(
+      result_a.begin(), result_b.begin(), thrust::make_transform_output_iterator(result_c.begin(), cuda::std::negate{})),
+    num_items,
+    plus_mul_neg{});
 
   // compute reference and verify
   c2h::host_vector<type> b_h = b;
-  c2h::host_vector<type> reference_h(num_items);
-  std::transform(a, a + num_items, b_h.begin(), reference_h.begin(), std::plus<type>{});
-  REQUIRE(reference_h == result);
+  c2h::host_vector<type> reference_a_h(num_items, thrust::no_init);
+  std::transform(a, a + num_items, b_h.begin(), reference_a_h.begin(), cuda::std::plus<type>{});
+  c2h::host_vector<type> reference_b_h(num_items, thrust::no_init);
+  std::transform(a, a + num_items, b_h.begin(), reference_b_h.begin(), cuda::std::multiplies<type>{});
+  CHECK(reference_a_h == result_a);
+  CHECK(reference_b_h == result_b);
+  CHECK(thrust::equal(a, a + num_items, result_c.begin()));
 }
 
 struct plus_needs_stable_address

cub/test/catch2_test_device_transform_api.cu

@@ -10,6 +10,37 @@
 
 #include <c2h/catch2_test_helper.h>
 
+// need a separate function because the ext. lambda needs to be enclosed by a function with external linkage on Windows
+void test_transform_many_many_api()
+{
+  // example-begin transform-many-many
+  auto input1 = thrust::device_vector<int>{0, -1, 2, -3, 4, -5};
+  auto input2 = thrust::device_vector<double>{5.2, 3.1, -1.1, 3.0, 3.2, 0.0};
+  auto op     = [] __device__(int a, double b) -> cuda::std::tuple<double, bool> {
+    const double product = a * b;
+    return {product, product < 0};
+  };
+
+  auto result1 = thrust::device_vector<double>(input1.size(), thrust::no_init);
+  auto result2 = thrust::device_vector<bool>(input1.size(), thrust::no_init);
+  cub::DeviceTransform::Transform(
+    cuda::std::tuple{input1.begin(), input2.begin()},
+    cuda::std::tuple{result1.begin(), result2.begin()},
+    input1.size(),
+    op);
+
+  const auto expected1 = thrust::host_vector<double>{0, -3.1, -2.2, -9, 12.8, -0};
+  const auto expected2 = thrust::host_vector<bool>{false, true, true, true, false, false};
+  // example-end transform-many-many
+  CHECK(result1 == expected1);
+  CHECK(result2 == expected2);
+}
+
+C2H_TEST("DeviceTransform::Transform many->many API example", "[device][device_transform]")
+{
+  test_transform_many_many_api();
+}
+
 // need a separate function because the ext. lambda needs to be enclosed by a function with external linkage on Windows
 void test_transform_api()
 {
@@ -21,7 +52,7 @@ void test_transform_api()
     return (a + b) * c;
   };
 
-  auto result = thrust::device_vector<int>(input1.size());
+  auto result = thrust::device_vector<int>(input1.size(), thrust::no_init);
   cub::DeviceTransform::Transform(
     cuda::std::tuple{input1.begin(), input2.begin(), input3}, result.begin(), input1.size(), op);
 
@@ -74,7 +105,7 @@ void test_transform_stable_api()
     return a + input2_ptr[i];
   };
 
-  auto result = thrust::device_vector<int>(input1.size());
+  auto result = thrust::device_vector<int>(input1.size(), thrust::no_init);
   cub::DeviceTransform::TransformStableArgumentAddresses(
     cuda::std::tuple{input1_ptr}, result.begin(), input1.size(), op);