Add static assert for task return and check for derived of ttg::corou…

…tine_handle TTs may return ttg::coroutine_handle (make_tt) or ttg::device::Task (custom TTs) Signed-off-by: Joseph Schuchart <[email protected]>
TESSEorg · Jun 4, 2024 · a42b682 · a42b682
1 parent 754c7d7
commit a42b682
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/ttg/ttg/tt.h b/ttg/ttg/tt.h
@@ -186,10 +186,15 @@ namespace ttg {
       id = ttg::TaskCoroutineID::Invalid;                                                                      \
     } else {                                                                                                   \
       auto coro_return = invoke;                                                                               \
-      if constexpr (std::is_same_v<decltype(coro_return), ttg::coroutine_handle<ttg::resumable_task_state>>)   \
+      static_assert(std::is_same_v<return_type, void> ||                                                         \
+                    std::is_base_of_v<ttg::coroutine_handle<ttg::resumable_task_state>, decltype(coro_return)>|| \
+                    std::is_base_of_v<ttg::coroutine_handle<ttg::device::detail::device_task_promise_type>,      \
+                                      decltype(coro_return)>);                                                   \
+      if constexpr (std::is_base_of_v<ttg::coroutine_handle<ttg::resumable_task_state>, decltype(coro_return)>)\
         id = ttg::TaskCoroutineID::ResumableTask;                                                              \
-      else if constexpr (std::is_same_v<decltype(coro_return),                                                 \
-                                        ttg::coroutine_handle<ttg::device::detail::device_task_promise_type>>) \
+      else if constexpr (std::is_base_of_v<                                                                    \
+                                        ttg::coroutine_handle<ttg::device::detail::device_task_promise_type>,  \
+                                        decltype(coro_return)>)                                                \
         id = ttg::TaskCoroutineID::DeviceTask;                                                                 \
       else                                                                                                     \
         std::abort();                                                                                          \