From c66708371a9042dbb80b2e5c6cd7c7d7cd66073e Mon Sep 17 00:00:00 2001
From: Cedric Augonnet <caugonnet@nvidia.com>
Date: Fri, 3 Jan 2025 11:30:19 +0100
Subject: [PATCH] Implement helpers to check if a function can be invoked from
 a tuple, or from a tuple where we removed tokens

---
 .../__stf/internal/backend_ctx.cuh            |  4 +-
 .../__stf/internal/void_interface.cuh         | 54 +++++++++++++++++++
 .../experimental/__stf/utility/traits.cuh     | 16 ++++++
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
index f0994095d21..15d05f1f894 100644
--- a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
+++ b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
@@ -197,9 +197,9 @@ public:
         delete w;
       };
 
-      constexpr bool fun_invocable_task_deps = ::std::is_invocable_v<Fun, decltype(payload)>;
+      constexpr bool fun_invocable_task_deps = reserved::is_tuple_invocable_v<Fun, decltype(payload)>;
       constexpr bool fun_invocable_task_non_void_deps =
-        reserved::is_invocable_with_filtered<Fun, decltype(payload)>::value;
+        reserved::is_tuple_invocable_with_filtered<Fun, decltype(payload)>::value;
 
       static_assert(fun_invocable_task_deps || fun_invocable_task_non_void_deps,
                     "Incorrect lambda function signature in host_launch.");
diff --git a/cudax/include/cuda/experimental/__stf/internal/void_interface.cuh b/cudax/include/cuda/experimental/__stf/internal/void_interface.cuh
index 8ea0993cb9f..17094894739 100644
--- a/cudax/include/cuda/experimental/__stf/internal/void_interface.cuh
+++ b/cudax/include/cuda/experimental/__stf/internal/void_interface.cuh
@@ -102,6 +102,25 @@ public:
 template <typename... Ts>
 using remove_void_interface_t = typename remove_void_interface<Ts...>::type;
 
+template <typename T>
+struct remove_void_interface_from_tuple
+{
+  // By default, if T is not a std::tuple, do nothing special
+  using type = T;
+};
+
+template <typename... Ts>
+struct remove_void_interface_from_tuple<::std::tuple<Ts...>>
+{
+  using type = remove_void_interface_t<Ts...>;
+};
+
+template <typename T>
+using remove_void_interface_from_tuple_t = typename remove_void_interface_from_tuple<T>::type;
+
+/**
+ * @brief Check if a function can be invoked while eliding arguments with a void_interface type.
+ */
 template <typename Fun, typename... Data>
 struct is_invocable_with_filtered
 {
@@ -130,6 +149,41 @@ public:
     decltype(check(::std::make_index_sequence<::std::tuple_size_v<remove_void_interface_t<Data...>>>{}))::value;
 };
 
+/**
+ * @brief Check if a function can be invoked using std::apply while eliding tuple arguments with a void_interface type.
+ */
+template <typename Fun, typename Tuple>
+struct is_tuple_invocable_with_filtered
+{
+private:
+  using filtered_tuple_t = remove_void_interface_from_tuple_t<Tuple>;
+
+  template <typename F, typename T>
+  static auto test(int) -> ::std::bool_constant<reserved::is_tuple_invocable_v<F, T>>
+  {
+    return {};
+  }
+
+  template <typename F>
+  static auto test(...) -> ::std::false_type
+  {
+    return {};
+  }
+
+  template <::std::size_t... Idx>
+  static auto check(::std::index_sequence<Idx...>)
+  {
+    return test<Fun, ::std::tuple_element_t<Idx, filtered_tuple_t>...>(0);
+  }
+
+public:
+  static constexpr bool value =
+    decltype(check(::std::make_index_sequence<::std::tuple_size_v<filtered_tuple_t>>{}))::value;
+};
+
+/**
+ * @brief Strip tuple entries with a "void_interface" type
+ */
 template <typename... Ts>
 auto remove_void_interface_types(const ::std::tuple<Ts...>& tpl)
 {
diff --git a/cudax/include/cuda/experimental/__stf/utility/traits.cuh b/cudax/include/cuda/experimental/__stf/utility/traits.cuh
index 8308e56d702..a30596f7bde 100644
--- a/cudax/include/cuda/experimental/__stf/utility/traits.cuh
+++ b/cudax/include/cuda/experimental/__stf/utility/traits.cuh
@@ -562,6 +562,22 @@ auto shuffled_array_tuple(ArgTypes... args)
 namespace reserved
 {
 
+/**
+ * @brief Trait class to check if a function can be invoked with std::apply using a tuple type
+ */
+template <typename F, typename Tuple>
+struct is_tuple_invocable : ::std::false_type
+{};
+
+// Partial specialization that unpacks the tuple
+template <typename F, typename... Args>
+struct is_tuple_invocable<F, ::std::tuple<Args...>> : ::std::is_invocable<F, Args...>
+{};
+
+// Convenient alias template
+template <typename F, typename Tuple>
+inline constexpr bool is_tuple_invocable_v = is_tuple_invocable<F, Tuple>::value;
+
 /**
  * @brief A compile-time boolean that checks if a type supports streaming with std::ostream <<.
  *