Implementation of task_group dynamic dependencies - part 1 - task_completion_handle

include/oneapi/tbb/detail/_task_handle.h

@@ -1,5 +1,6 @@
 /*
-    Copyright (c) 2020-2024 Intel Corporation
+    Copyright (c) 2020-2025 Intel Corporation
+    Copyright (c) 2025 UXL Foundation Contributors
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -32,11 +33,39 @@ namespace d2 {
 
 class task_handle;
 
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+
+class task_dynamic_state {
+public:
+    task_dynamic_state(d1::small_object_allocator& alloc)
+        : m_num_references(1) // reserves a task co-ownership for dynamic state
+        , m_allocator(alloc)
+    {}
+
+    void reserve() { ++m_num_references; }
+
+    void release() {
+        if (--m_num_references == 0) {
+            m_allocator.delete_object(this);
+        }
+    }
+
+    void complete_task() {
+    }
+private:
+    std::atomic<std::size_t> m_num_references;
+    d1::small_object_allocator m_allocator;
+};
+#endif // __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+
 class task_handle_task : public d1::task {
     std::uint64_t m_version_and_traits{};
     d1::wait_tree_vertex_interface* m_wait_tree_vertex;
     d1::task_group_context& m_ctx;
     d1::small_object_allocator m_allocator;
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+    std::atomic<task_dynamic_state*> m_dynamic_state;
+#endif
 public:
     void finalize(const d1::execution_data* ed = nullptr) {
         if (ed) {
@@ -49,16 +78,56 @@ class task_handle_task : public d1::task {
     task_handle_task(d1::wait_tree_vertex_interface* vertex, d1::task_group_context& ctx, d1::small_object_allocator& alloc)
         : m_wait_tree_vertex(vertex)
         , m_ctx(ctx)
-        , m_allocator(alloc) {
+        , m_allocator(alloc)
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+        , m_dynamic_state(nullptr)
+#endif
+    {
         suppress_unused_warning(m_version_and_traits);
         m_wait_tree_vertex->reserve();
     }
 
     ~task_handle_task() override {
         m_wait_tree_vertex->release();
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+        task_dynamic_state* current_state = m_dynamic_state.load(std::memory_order_relaxed);
+        if (current_state != nullptr) {
+            current_state->release();
+        }
+#endif
     }
 
     d1::task_group_context& ctx() const { return m_ctx; }
+
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+    // Returns the dynamic state associated with the task. If the state has not been initialized, initializes it.
+    task_dynamic_state* get_dynamic_state() {
+        task_dynamic_state* current_state = m_dynamic_state.load(std::memory_order_acquire);
+
+        if (current_state == nullptr) {
+            d1::small_object_allocator alloc;
+
+            task_dynamic_state* new_state = alloc.new_object<task_dynamic_state>(alloc);
+
+            if (m_dynamic_state.compare_exchange_strong(current_state, new_state)) {
+                current_state = new_state;
+            } else {
+                // CAS failed, current_state points to the dynamic state created by another thread
+                alloc.delete_object(new_state);
+            }
+        }
+
+        __TBB_ASSERT(current_state != nullptr, "Failed to create dynamic state");
+        return current_state;
+    }
+
+    void complete_task() {
+        task_dynamic_state* current_state = m_dynamic_state.load(std::memory_order_relaxed);
+        if (current_state != nullptr) {
+            current_state->complete_task();
+        }
+    }
+#endif
 };
 
 
@@ -84,21 +153,26 @@ class task_handle {
 
 private:
     friend struct task_handle_accessor;
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+    friend class task_completion_handle;
+#endif
 
-    task_handle(task_handle_task* t) : m_handle {t}{};
+    task_handle(task_handle_task* t) : m_handle {t}{}
 
     d1::task* release() {
        return m_handle.release();
     }
 };
 
 struct task_handle_accessor {
-static task_handle              construct(task_handle_task* t)  { return {t}; }
-static d1::task*                release(task_handle& th)        { return th.release(); }
-static d1::task_group_context&  ctx_of(task_handle& th)         {
-    __TBB_ASSERT(th.m_handle, "ctx_of does not expect empty task_handle.");
-    return th.m_handle->ctx();
-}
+    static task_handle construct(task_handle_task* t) { return {t}; }
+
+    static d1::task* release(task_handle& th) { return th.release(); }
+
+    static d1::task_group_context& ctx_of(task_handle& th) {
+        __TBB_ASSERT(th.m_handle, "ctx_of does not expect empty task_handle.");
+        return th.m_handle->ctx();
+    }
 };
 
 inline bool operator==(task_handle const& th, std::nullptr_t) noexcept {
@@ -116,6 +190,108 @@ inline bool operator!=(std::nullptr_t, task_handle const& th) noexcept {
     return th.m_handle != nullptr;
 }
 
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+class task_completion_handle {
+public:
+    task_completion_handle() : m_task_state(nullptr) {}
+
+    task_completion_handle(const task_completion_handle& other) 
+        : m_task_state(other.m_task_state)
+    {
+        // Register one more co-owner of the dynamic state
+        if (m_task_state) m_task_state->reserve();
+    }
+    task_completion_handle(task_completion_handle&& other)
+        : m_task_state(other.m_task_state)
+    {
+        other.m_task_state = nullptr;
+    }
+
+    task_completion_handle(const task_handle& th)
+        : m_task_state(nullptr)
+    {
+        __TBB_ASSERT(th, "Construction of task_completion_handle from an empty task_handle");
+        m_task_state = th.m_handle->get_dynamic_state();
+        // Register one more co-owner of the dynamic state
+        m_task_state->reserve();
+    }
+
+    ~task_completion_handle() {
+        if (m_task_state) m_task_state->release();
+    }
+
+    task_completion_handle& operator=(const task_completion_handle& other) {
+        if (m_task_state != other.m_task_state) {
+            // Release co-ownership on the previously tracked dynamic state
+            if (m_task_state) m_task_state->release();
+
+            m_task_state = other.m_task_state;
+
+            // Register new co-owner of the new dynamic state
+            if (m_task_state) m_task_state->reserve();
+        }
+        return *this;
+    }
+
+    task_completion_handle& operator=(task_completion_handle&& other) {
+        if (this != &other) {
+            // Release co-ownership on the previously tracked dynamic state
+            if (m_task_state) m_task_state->release();
+
+            m_task_state = other.m_task_state;
+            other.m_task_state = nullptr;
+        }
+        return *this;
+    }
+
+    task_completion_handle& operator=(const task_handle& th) {
+        __TBB_ASSERT(th, "Assignment of task_completion_state from an empty task_handle");
+        task_dynamic_state* th_state = th.m_handle->get_dynamic_state();
+        __TBB_ASSERT(th_state != nullptr, "No state in the non-empty task_handle");
+        if (m_task_state != th_state) {
+            // Release co-ownership on the previously tracked dynamic state
+            if (m_task_state) m_task_state->release();
+
+            m_task_state = th_state;
+
+            // Reserve co-ownership on the new dynamic state
+            m_task_state->reserve();
+        }
+        return *this;
+    }
+
+    explicit operator bool() const noexcept { return m_task_state != nullptr; }
+private:
+    friend bool operator==(const task_completion_handle& t, std::nullptr_t) noexcept {
+        return t.m_task_state == nullptr;
+    }
+
+    friend bool operator==(const task_completion_handle& lhs, const task_completion_handle& rhs) noexcept {
+        return lhs.m_task_state == rhs.m_task_state;
+    }
+
+#if !__TBB_CPP20_COMPARISONS_PRESENT
+    friend bool operator==(std::nullptr_t, const task_completion_handle& t) noexcept {
+        return t == nullptr;
+    }
+
+    friend bool operator!=(const task_completion_handle& t, std::nullptr_t) noexcept {
+        return !(t == nullptr);
+    }
+
+    friend bool operator!=(std::nullptr_t, const task_completion_handle& t) noexcept {
+        return !(t == nullptr);
+    }
+
+    friend bool operator!=(const task_completion_handle& lhs, const task_completion_handle& rhs) noexcept {
+        return !(lhs == rhs);
+    }
+#endif // !__TBB_CPP20_COMPARISONS_PRESENT
+
+    task_dynamic_state* m_task_state;
+};
+#endif
+
 } // namespace d2
 } // namespace detail
 } // namespace tbb

include/oneapi/tbb/task_group.h

@@ -1,5 +1,6 @@
 /*
-    Copyright (c) 2005-2024 Intel Corporation
+    Copyright (c) 2005-2025 Intel Corporation
+    Copyright (c) 2025 UXL Foundation Contributors
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -87,6 +88,9 @@ class function_task : public task_handle_task  {
     d1::task* execute(d1::execution_data& ed) override {
         __TBB_ASSERT(ed.context == &this->ctx(), "The task group context should be used for all tasks");
         task* res = task_ptr_or_nullptr(m_func);
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+        this->complete_task();
+#endif
         finalize(&ed);
         return res;
     }
@@ -701,6 +705,9 @@ using detail::d1::is_current_task_group_canceling;
 using detail::r1::missing_wait;
 
 using detail::d2::task_handle;
+#if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
+using detail::d2::task_completion_handle;
+#endif
 }
 
 } // namespace tbb