[CUDAX] Add forwarding reference to functor accepting launch (#2677)

* Add forwarding reference to functor accepting launch * Add missing forward
NVIDIA · Nov 1, 2024 · 0b02ae7 · 0b02ae7
1 parent 7d75492
commit 0b02ae7
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/cudax/include/cuda/experimental/__launch/launch.cuh b/cudax/include/cuda/experimental/__launch/launch.cuh
@@ -119,7 +119,7 @@ launch_impl(::cuda::stream_ref stream, Config conf, const Kernel& kernel_fn, con
  */
 template <typename... Args, typename... Config, typename Dimensions, typename Kernel>
 void launch(
-  ::cuda::stream_ref stream, const kernel_config<Dimensions, Config...>& conf, const Kernel& kernel, Args... args)
+  ::cuda::stream_ref stream, const kernel_config<Dimensions, Config...>& conf, const Kernel& kernel, Args&&... args)
 {
   __ensure_current_device __dev_setter(stream);
   cudaError_t status;
@@ -132,14 +132,18 @@ void launch(
       launcher,
       conf,
       kernel,
-      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, args))...);
+      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, std::forward<Args>(args)))...);
   }
   else
   {
     static_assert(::cuda::std::is_invocable_v<Kernel, as_kernel_arg_t<Args>...>);
     auto launcher = detail::kernel_launcher_no_config<Kernel, as_kernel_arg_t<Args>...>;
     status        = detail::launch_impl(
-      stream, conf, launcher, kernel, static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, args))...);
+      stream,
+      conf,
+      launcher,
+      kernel,
+      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, std::forward<Args>(args)))...);
   }
   if (status != cudaSuccess)
   {
@@ -189,7 +193,7 @@ void launch(
  * arguments to be passed into the kernel functor
  */
 template <typename... Args, typename... Levels, typename Kernel>
-void launch(::cuda::stream_ref stream, const hierarchy_dimensions<Levels...>& dims, const Kernel& kernel, Args... args)
+void launch(::cuda::stream_ref stream, const hierarchy_dimensions<Levels...>& dims, const Kernel& kernel, Args&&... args)
 {
   __ensure_current_device __dev_setter(stream);
   cudaError_t status;
@@ -202,7 +206,7 @@ void launch(::cuda::stream_ref stream, const hierarchy_dimensions<Levels...>& di
       launcher,
       dims,
       kernel,
-      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, args))...);
+      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, std::forward<Args>(args)))...);
   }
   else
   {
@@ -213,7 +217,7 @@ void launch(::cuda::stream_ref stream, const hierarchy_dimensions<Levels...>& di
       kernel_config(dims),
       launcher,
       kernel,
-      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, args))...);
+      static_cast<as_kernel_arg_t<Args>>(detail::__launch_transform(stream, std::forward<Args>(args)))...);
   }
   if (status != cudaSuccess)
   {