Add test case for empty device::select() and fix bug in make_tt

The code path for no-arg callables was missing a return and we have to specify template arguments for std::array explicitly because the array could be empty. Signed-off-by: Joseph Schuchart <[email protected]>
devreal · Nov 18, 2024 · 80c25b2 · 80c25b2
1 parent b932ac0
commit 80c25b2
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/tests/unit/device_coro.cc b/tests/unit/device_coro.cc
@@ -426,13 +426,26 @@ TEST_CASE("Device", "coro") {
 #endif // TTG_HAVE_CUDA
       }
     };
-
     auto tt = ttg::make_tt<ttg::ExecutionSpace::CUDA>(fn, ttg::edges(edge), ttg::edges(edge),
                                                       "device_task", {"edge_in"}, {"edge_out"});
     ttg::make_graph_executable(tt);
     if (ttg::default_execution_context().rank() == 0) tt->invoke(0, value_t{});
     ttg::ttg_fence(ttg::default_execution_context());
   }
+
+  SECTION("empty-select") {
+    ttg::Edge<void, void> edge;
+    auto fn = []() -> ttg::device::Task {
+      co_await ttg::device::select();
+      /* nothing else to do */
+    };
+    auto tt = ttg::make_tt<ttg::ExecutionSpace::CUDA>(fn, ttg::edges(edge), ttg::edges(),
+                                                      "device_task", {"edge_in"}, {"edge_out"});
+    ttg::make_graph_executable(tt);
+    tt->invoke();
+    ttg::ttg_fence(ttg::default_execution_context());
+  };
+
 }
 
 #endif // TTG_IMPL_DEVICE_SUPPORT
diff --git a/ttg/ttg/device/task.h b/ttg/ttg/device/task.h
@@ -37,7 +37,7 @@ namespace ttg::device {
     template<typename... Ts, std::size_t... Is>
     auto extract_buffer_data(detail::to_device_t<Ts...>& a, std::index_sequence<Is...>) {
       using arg_types = std::tuple<Ts...>;
-      return std::array{
+      return std::array<device_input_data_t, sizeof...(Ts)>{
                 device_input_data_t{TTG_IMPL_NS::buffer_data(std::get<Is>(a.ties)),
                                     std::get<Is>(a.ties).scope(),
                                     ttg::meta::is_const_v<std::tuple_element_t<Is, arg_types>>,

diff --git a/ttg/ttg/make_tt.h b/ttg/ttg/make_tt.h
@@ -213,9 +213,9 @@ class CallableWrapTT
 
     auto invoke_func_empty_tuple = [&](auto&&... args){
       if constexpr(funcT_receives_input_tuple) {
-        invoke_func_handle_ret(std::tuple<>{}, std::forward<decltype(args)>(args)...);
+        return invoke_func_handle_ret(std::tuple<>{}, std::forward<decltype(args)>(args)...);
       } else {
-        invoke_func_handle_ret(std::forward<decltype(args)>(args)...);
+        return invoke_func_handle_ret(std::forward<decltype(args)>(args)...);
       }
     };