Rework acquisition of the current task

Signed-off-by: pavelkumbrasev <[email protected]>

include/oneapi/tbb/detail/_task.h

@@ -61,7 +61,9 @@ using suspend_callback_type = void(*)(void*, suspend_point_type*);
 TBB_EXPORT void __TBB_EXPORTED_FUNC suspend(suspend_callback_type suspend_callback, void* user_callback);
 TBB_EXPORT void __TBB_EXPORTED_FUNC resume(suspend_point_type* tag);
 TBB_EXPORT suspend_point_type* __TBB_EXPORTED_FUNC current_suspend_point();
+
 TBB_EXPORT void __TBB_EXPORTED_FUNC notify_waiters(std::uintptr_t wait_ctx_addr);
+TBB_EXPORT d1::task* __TBB_EXPORTED_FUNC current_task();
 
 class thread_data;
 class task_dispatcher;
@@ -199,6 +201,10 @@ inline void wait(wait_context& wait_ctx, task_group_context& ctx) {
     call_itt_task_notify(destroy, &wait_ctx);
 }
 
+inline d1::task* current_task() {
+    return r1::current_task();
+}
+
 using r1::current_context;
 
 class task_traits {

include/oneapi/tbb/task_group.h

@@ -69,10 +69,6 @@ TBB_EXPORT bool __TBB_EXPORTED_FUNC cancel_group_execution(d1::task_group_contex
 TBB_EXPORT bool __TBB_EXPORTED_FUNC is_group_execution_cancelled(d1::task_group_context&);
 TBB_EXPORT void __TBB_EXPORTED_FUNC capture_fp_settings(d1::task_group_context&);
 
-TBB_EXPORT void __TBB_EXPORTED_FUNC set_top_group_task(d1::task*);
-TBB_EXPORT d1::task* __TBB_EXPORTED_FUNC get_top_group_task();
-TBB_EXPORT void __TBB_EXPORTED_FUNC unset_top_group_task();
-
 struct task_group_context_impl;
 }
 
@@ -527,15 +523,6 @@ class function_task : public base_task_group_task {
     }
 
     task* execute(execution_data& ed) override {
-        struct top_task_guard_type {
-            top_task_guard_type(d1::task* t) {
-                r1::set_top_group_task(t);
-            }
-            ~top_task_guard_type() {
-                r1::unset_top_group_task();
-            }
-        } top_task_guard{this};
-
         task* res = d2::task_ptr_or_nullptr(m_func);
         finalize(ed);
         return res;
@@ -614,7 +601,7 @@ class task_group_base : no_copy {
 
     template<typename F>
     task* prepare_task(F&& f) {
-        base_task_group_task* parent_task = static_cast<base_task_group_task*>(r1::get_top_group_task());
+        base_task_group_task* parent_task = dynamic_cast<base_task_group_task*>(current_task());
         task_group_continuation* continuation = nullptr;
 
         if (parent_task && parent_task->is_same_task_group(&m_wait_ctx)) {

src/tbb/def/lin32-tbb.def

@@ -77,6 +77,7 @@ _ZN3tbb6detail2r17suspendEPFvPvPNS1_18suspend_point_typeEES2_;
 _ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE;
 _ZN3tbb6detail2r121current_suspend_pointEv;
 _ZN3tbb6detail2r114notify_waitersEj;
+_ZN3tbb6detail2r112current_taskEv;
 
 /* Task dispatcher (task_dispatcher.cpp) */
 _ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE;
@@ -95,11 +96,6 @@ _ZN3tbb6detail2r15resetERNS0_2d118task_group_contextE;
 _ZN3tbb6detail2r17destroyERNS0_2d118task_group_contextE;
 _ZN3tbb6detail2r119capture_fp_settingsERNS0_2d118task_group_contextE;
 
-/* Task group (task_group_context.cpp) */
-_ZN3tbb6detail2r118set_top_group_taskEPNS0_2d14taskE;
-_ZN3tbb6detail2r118get_top_group_taskEv;
-_ZN3tbb6detail2r120unset_top_group_taskEv;
-
 /* Task arena (arena.cpp) */
 _ZN3tbb6detail2r115max_concurrencyEPKNS0_2d115task_arena_baseE;
 _ZN3tbb6detail2r110initializeERNS0_2d115task_arena_baseE;

src/tbb/def/lin64-tbb.def

@@ -77,6 +77,7 @@ _ZN3tbb6detail2r17suspendEPFvPvPNS1_18suspend_point_typeEES2_;
 _ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE;
 _ZN3tbb6detail2r121current_suspend_pointEv;
 _ZN3tbb6detail2r114notify_waitersEm;
+_ZN3tbb6detail2r112current_taskEv;
 
 /* Task dispatcher (task_dispatcher.cpp) */
 _ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE;
@@ -95,11 +96,6 @@ _ZN3tbb6detail2r15resetERNS0_2d118task_group_contextE;
 _ZN3tbb6detail2r17destroyERNS0_2d118task_group_contextE;
 _ZN3tbb6detail2r119capture_fp_settingsERNS0_2d118task_group_contextE;
 
-/* Task group (task_group_context.cpp) */
-_ZN3tbb6detail2r118set_top_group_taskEPNS0_2d14taskE;
-_ZN3tbb6detail2r118get_top_group_taskEv;
-_ZN3tbb6detail2r120unset_top_group_taskEv;
-
 /* Task arena (arena.cpp) */
 _ZN3tbb6detail2r115max_concurrencyEPKNS0_2d115task_arena_baseE;
 _ZN3tbb6detail2r110initializeERNS0_2d115task_arena_baseE;

src/tbb/def/mac64-tbb.def

@@ -79,6 +79,7 @@ __ZN3tbb6detail2r17suspendEPFvPvPNS1_18suspend_point_typeEES2_
 __ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE
 __ZN3tbb6detail2r121current_suspend_pointEv
 __ZN3tbb6detail2r114notify_waitersEm
+__ZN3tbb6detail2r112current_taskEv
 
 # Task dispatcher (task_dispatcher.cpp)
 __ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE
@@ -97,11 +98,6 @@ __ZN3tbb6detail2r15resetERNS0_2d118task_group_contextE
 __ZN3tbb6detail2r17destroyERNS0_2d118task_group_contextE
 __ZN3tbb6detail2r119capture_fp_settingsERNS0_2d118task_group_contextE
 
-# Task group (task_group_context.cpp)
-__ZN3tbb6detail2r118set_top_group_taskEPNS0_2d14taskE
-__ZN3tbb6detail2r118get_top_group_taskEv
-__ZN3tbb6detail2r120unset_top_group_taskEv
-
 # Task arena (arena.cpp)
 __ZN3tbb6detail2r115max_concurrencyEPKNS0_2d115task_arena_baseE
 __ZN3tbb6detail2r110initializeERNS0_2d115task_arena_baseE

src/tbb/def/win32-tbb.def

@@ -71,6 +71,7 @@ EXPORTS
 ?resume@r1@detail@tbb@@YAXPAUsuspend_point_type@123@@Z
 ?suspend@r1@detail@tbb@@YAXP6AXPAXPAUsuspend_point_type@123@@Z0@Z
 ?notify_waiters@r1@detail@tbb@@YAXI@Z
+?current_task@r1@detail@tbb@@YAPAVtask@d1@23@XZ
 
 ; Task dispatcher (task_dispatcher.cpp)
 ?spawn@r1@detail@tbb@@YAXAAVtask@d1@23@AAVtask_group_context@523@G@Z
@@ -89,11 +90,6 @@ EXPORTS
 ?is_group_execution_cancelled@r1@detail@tbb@@YA_NAAVtask_group_context@d1@23@@Z
 ?reset@r1@detail@tbb@@YAXAAVtask_group_context@d1@23@@Z
 
-; Task group (task_group_context.cpp)
-?set_top_group_task@r1@detail@tbb@@YAXPAVtask@d1@23@@Z
-?get_top_group_task@r1@detail@tbb@@YAPAVtask@d1@23@XZ
-?unset_top_group_task@r1@detail@tbb@@YAXXZ
-
 ; Task arena (arena.cpp)
 ?attach@r1@detail@tbb@@YA_NAAVtask_arena_base@d1@23@@Z
 ?enqueue@r1@detail@tbb@@YAXAAVtask@d1@23@PAVtask_arena_base@523@@Z

src/tbb/def/win64-tbb.def

@@ -71,6 +71,7 @@ EXPORTS
 ?resume@r1@detail@tbb@@YAXPEAUsuspend_point_type@123@@Z
 ?current_suspend_point@r1@detail@tbb@@YAPEAUsuspend_point_type@123@XZ
 ?notify_waiters@r1@detail@tbb@@YAX_K@Z
+?current_task@r1@detail@tbb@@YAPEAVtask@d1@23@XZ
 
 ; Task dispatcher (task_dispatcher.cpp)
 ?spawn@r1@detail@tbb@@YAXAEAVtask@d1@23@AEAVtask_group_context@523@@Z
@@ -89,11 +90,6 @@ EXPORTS
 ?cancel_group_execution@r1@detail@tbb@@YA_NAEAVtask_group_context@d1@23@@Z
 ?capture_fp_settings@r1@detail@tbb@@YAXAEAVtask_group_context@d1@23@@Z
 
-; Task group (task_group_context.cpp)
-?set_top_group_task@r1@detail@tbb@@YAXPEAVtask@d1@23@@Z
-?get_top_group_task@r1@detail@tbb@@YAPEAVtask@d1@23@XZ
-?unset_top_group_task@r1@detail@tbb@@YAXXZ
-
 ; Task arena (arena.cpp)
 ?max_concurrency@r1@detail@tbb@@YAHPEBVtask_arena_base@d1@23@@Z
 ?initialize@r1@detail@tbb@@YAXAEAVtask_arena_base@d1@23@@Z

src/tbb/scheduler_common.h

@@ -41,7 +41,6 @@
 
 #include <cstdint>
 #include <exception>
-#include <stack>
 #include <memory> // unique_ptr
 
 //! Mutex type for global locks in the scheduler
@@ -480,8 +479,8 @@ class alignas (max_nfs_size) task_dispatcher {
     //! Suspend point (null if this task dispatcher has been never suspended)
     suspend_point_type* m_suspend_point{ nullptr };
 
-    //! Used to build tree-like structure in task_group
-    std::stack<d1::task*> my_task_group_tasks;
+    //! Innermost task whose task::execute() is running. A nullptr on the outermost level.
+    d1::task* m_innermost_running_task{ nullptr };
 
     //! Attempt to get a task from the mailbox.
     /** Gets a task only if it has not been executed by its sender or a thief

src/tbb/task.cpp

@@ -221,6 +221,10 @@ void notify_waiters(std::uintptr_t wait_ctx_addr) {
     governor::get_thread_data()->my_arena->get_waiting_threads_monitor().notify(is_related_wait_ctx);
 }
 
+d1::task* current_task() {
+    return governor::get_thread_data()->get_current_task();
+}
+
 } // namespace r1
 } // namespace detail
 } // namespace tbb

src/tbb/task_dispatcher.h

@@ -249,15 +249,17 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
         task_dispatcher& task_disp;
         execution_data_ext old_execute_data_ext;
         properties old_properties;
+        d1::task* old_innermost_running_task;
 
         ~dispatch_loop_guard() {
             task_disp.m_execute_data_ext = old_execute_data_ext;
             task_disp.m_properties = old_properties;
+            task_disp.m_innermost_running_task = old_innermost_running_task;
 
             __TBB_ASSERT(task_disp.m_thread_data && governor::is_thread_data_set(task_disp.m_thread_data), nullptr);
             __TBB_ASSERT(task_disp.m_thread_data->my_task_dispatcher == &task_disp, nullptr);
         }
-    } dl_guard{ *this, m_execute_data_ext, m_properties };
+    } dl_guard{ *this, m_execute_data_ext, m_properties, m_innermost_running_task };
 
     // The context guard to track fp setting and itt tasks.
     context_guard_helper</*report_tasks=*/ITTPossible> context_guard;
@@ -317,6 +319,7 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
 
                     ITT_CALLEE_ENTER(ITTPossible, t, itt_caller);
 
+                    m_innermost_running_task = t;
                     if (ed.context->is_group_execution_cancelled()) {
                         t = t->cancel(ed);
                     } else {

src/tbb/task_group_context.cpp

@@ -1,5 +1,5 @@
 /*
-    Copyright (c) 2005-2024 Intel Corporation
+    Copyright (c) 2005-2023 Intel Corporation
 
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
@@ -333,26 +333,6 @@ void task_group_context_impl::copy_fp_settings(d1::task_group_context& ctx, cons
     However this version of the algorithm requires more analysis and verification.
 */
 
-//------------------------------------------------------------------------
-// task_group
-//------------------------------------------------------------------------
-
-void __TBB_EXPORTED_FUNC set_top_group_task(d1::task* t) {
-    governor::get_thread_data()->set_top_group_task(t);
-}
-
-d1::task* __TBB_EXPORTED_FUNC get_top_group_task() {
-    return governor::get_thread_data()->get_top_group_task();
-}
-
-void __TBB_EXPORTED_FUNC unset_top_group_task() {
-    governor::get_thread_data()->unset_top_group_task();
-}
-
-//------------------------------------------------------------------------
-// task_group_context
-//------------------------------------------------------------------------
-
 void __TBB_EXPORTED_FUNC initialize(d1::task_group_context& ctx) {
     task_group_context_impl::initialize(ctx);
 }

src/tbb/thread_data.h

@@ -139,9 +139,7 @@ class thread_data : public ::rml::job
     void leave_task_dispatcher();
     void propagate_task_group_state(std::atomic<uint32_t> d1::task_group_context::* mptr_state, d1::task_group_context& src, uint32_t new_state);
 
-    void set_top_group_task(d1::task* t);
-    d1::task* get_top_group_task();
-    void unset_top_group_task();
+    d1::task* get_current_task();
 
     //! Index of the arena slot the scheduler occupies now, or occupied last time
     unsigned short my_arena_index;
@@ -255,17 +253,8 @@ inline void thread_data::propagate_task_group_state(std::atomic<std::uint32_t> d
     my_context_list->epoch.store(the_context_state_propagation_epoch.load(std::memory_order_relaxed), std::memory_order_release);
 }
 
-inline void thread_data::set_top_group_task(d1::task* t) {
-    my_task_dispatcher->my_task_group_tasks.push(t);
-}
-
-inline d1::task* thread_data::get_top_group_task() {
-    return my_task_dispatcher->my_task_group_tasks.empty() ? nullptr : my_task_dispatcher->my_task_group_tasks.top();
-}
-
-inline void thread_data::unset_top_group_task() {
-    __TBB_ASSERT(!my_task_dispatcher->my_task_group_tasks.empty(), nullptr);
-    my_task_dispatcher->my_task_group_tasks.pop();
+inline d1::task* thread_data::get_current_task() {
+    return my_task_dispatcher->m_innermost_running_task;
 }
 
 } // namespace r1