Simplify cuda::host_launch API

Author: davebayer

libcudacxx/include/cuda/__launch/host_launch.h

@@ -24,15 +24,24 @@
 #include <cuda/__driver/driver_api.h>
 #include <cuda/__stream/stream_ref.h>
 #include <cuda/std/__functional/reference_wrapper.h>
+#include <cuda/std/__tuple_dir/apply.h>
+#include <cuda/std/__tuple_dir/tuple.h>
 #include <cuda/std/__type_traits/decay.h>
+#include <cuda/std/__type_traits/is_function.h>
 #include <cuda/std/__type_traits/is_move_constructible.h>
-#include <cuda/std/__utility/forward.h>
-#include <cuda/std/tuple>
+#include <cuda/std/__type_traits/is_pointer.h>
+#include <cuda/std/__utility/move.h>
 
 #include <cuda/std/__cccl/prologue.h>
 
 _CCCL_BEGIN_NAMESPACE_CUDA
 
+template <class _Callable>
+_CCCL_HOST_API inline void CUDA_CB __host_func_launcher(void* __callable_ptr)
+{
+  (*static_cast<_Callable*>(__callable_ptr))();
+}
+
 template <class _Callable, class... _Args>
 struct __stream_callback_data
 {
@@ -73,37 +82,29 @@ _CCCL_HOST_API void host_launch(stream_ref __stream, _Callable __callable, _Args
   static_assert((::cuda::std::is_move_constructible_v<_Args> && ...),
                 "All callback arguments must be move constructible");
 
-  using _CallbackData                = __stream_callback_data<_Callable, _Args...>;
-  _CallbackData* __callback_data_ptr = new _CallbackData{::cuda::std::move(__callable), {::cuda::std::move(__args)...}};
+  constexpr auto __has_args = sizeof...(_Args) > 0;
 
-  // We use the callback here to have it execute even on stream error, because it needs to free the above allocation
-  ::cuda::__driver::__streamAddCallback(__stream.get(), __stream_callback_launcher<_CallbackData>, __callback_data_ptr);
-}
+  if constexpr (!__has_args && ::cuda::std::is_function_v<_Callable> && ::cuda::std::is_pointer_v<_Callable>)
+  {
+    ::cuda::__driver::__launchHostFunc(__stream.get(), ::cuda::__host_func_launcher<_Callable>, __callable);
+  }
+  else if constexpr (!__has_args && ::cuda::std::__is_cuda_std_reference_wrapper_v<_Callable>)
+  {
+    ::cuda::__driver::__launchHostFunc(
+      __stream.get(), ::cuda::__host_func_launcher<typename _Callable::type>, ::cuda::std::addressof(__callable.get()));
+  }
+  else
+  {
+    using _CallbackData = __stream_callback_data<_Callable, _Args...>;
+    _CallbackData* __callback_data_ptr =
+      new _CallbackData{::cuda::std::move(__callable), {::cuda::std::move(__args)...}};
 
-template <class _Callable>
-_CCCL_HOST_API inline void CUDA_CB __host_func_launcher(void* __callable_ptr)
-{
-  (*static_cast<_Callable*>(__callable_ptr))();
+    // We use the callback here to have it execute even on stream error, because it needs to free the above allocation
+    ::cuda::__driver::__streamAddCallback(
+      __stream.get(), ::cuda::__stream_callback_launcher<_CallbackData>, __callback_data_ptr);
+  }
 }
 
-//! @brief Launches a host callable to be executed in stream order on the provided stream
-//!
-//! Callable will be called using the supplied reference. If the callable was destroyed
-//! or moved by the time it is asynchronously called the behavior is undefined.
-//!
-//! Callable must not call any APIs from cuda, thrust or cub namespaces.
-//! It must not call into CUDA Runtime or Driver APIs. It also can't depend on another
-//! thread that might block on any asynchronous CUDA work.
-//!
-//! @param __stream Stream to launch the host function on
-//! @param __callable A reference to a host function or callable object to call in stream order
-template <class _Callable>
-_CCCL_HOST_API void host_launch(stream_ref __stream, ::cuda::std::reference_wrapper<_Callable> __callable)
-{
-  static_assert(::cuda::std::is_invocable_v<_Callable>, "Callable in reference_wrapper can't take any arguments");
-  ::cuda::__driver::__launchHostFunc(
-    __stream.get(), __host_func_launcher<_Callable>, ::cuda::std::addressof(__callable.get()));
-}
 _CCCL_END_NAMESPACE_CUDA
 
 #include <cuda/std/__cccl/epilogue.h>

libcudacxx/test/libcudacxx/cuda/ccclrt/launch/host_launch.cu

@@ -32,6 +32,21 @@ void unblock_and_wait_stream(cuda::stream_ref stream, cuda::atomic<int>& atomic)
   atomic = 0;
 }
 
+bool ordinary_function_run_proof = false;
+
+template <class Ret, class... Args>
+Ret ordinary_function(Args...)
+{
+  ordinary_function_run_proof = true;
+  return (Ret) 0;
+}
+
+[[nodiscard]] int nodiscard_ordinary_function()
+{
+  ordinary_function_run_proof = true;
+  return 0;
+}
+
 void launch_local_lambda(cuda::stream_ref stream, int& set, int set_to)
 {
   auto lambda = [&set, set_to]() {
@@ -111,9 +126,51 @@ C2H_CCCLRT_TEST("Host launch", "")
   cuda::device_ref device{0};
   device.init();
 
-  cuda::atomic<int> atomic = 0;
   cuda::stream stream{device};
-  int i = 0;
+
+  SECTION("Ordinary function without arguments returning void")
+  {
+    CCCLRT_REQUIRE(ordinary_function_run_proof == false);
+
+    cuda::host_launch(stream, ordinary_function<void>);
+
+    stream.sync();
+    CCCLRT_REQUIRE(ordinary_function_run_proof == true);
+    ordinary_function_run_proof = false;
+  }
+  SECTION("Ordinary function without arguments returning int")
+  {
+    CCCLRT_REQUIRE(ordinary_function_run_proof == false);
+
+    cuda::host_launch(stream, ordinary_function<int>);
+
+    stream.sync();
+    CCCLRT_REQUIRE(ordinary_function_run_proof == true);
+    ordinary_function_run_proof = false;
+  }
+  SECTION("Ordinary function with arguments returning void")
+  {
+    CCCLRT_REQUIRE(ordinary_function_run_proof == false);
+
+    cuda::host_launch(stream, ordinary_function<int, char, double>, 'c', 1.0);
+
+    stream.sync();
+    CCCLRT_REQUIRE(ordinary_function_run_proof == true);
+    ordinary_function_run_proof = false;
+  }
+  SECTION("Nodiscard ordinary function")
+  {
+    CCCLRT_REQUIRE(ordinary_function_run_proof == false);
+
+    cuda::host_launch(stream, nodiscard_ordinary_function);
+
+    stream.sync();
+    CCCLRT_REQUIRE(ordinary_function_run_proof == true);
+    ordinary_function_run_proof = false;
+  }
+
+  cuda::atomic<int> atomic = 0;
+  int i                    = 0;
 
   auto set_lambda = [&](int set) {
     i = set;