Add current_task_ptr entry point (#1821)

kboyarinov · akukanov · web-flow · commit 575948b37cec · 2025-08-27T19:21:31.000+03:00
Co-authored-by: Alexey Kukanov &lt;alexey.kukanov@intel.com&gt;
diff --git a/include/oneapi/tbb/detail/_task.h b/include/oneapi/tbb/detail/_task.h
@@ -63,6 +63,7 @@ TBB_EXPORT d1::slot_id __TBB_EXPORTED_FUNC execution_slot(const d1::execution_da
 TBB_EXPORT d1::slot_id __TBB_EXPORTED_FUNC execution_slot(const d1::task_arena_base&);
 TBB_EXPORT d1::task_group_context* __TBB_EXPORTED_FUNC current_context();
 TBB_EXPORT d1::wait_tree_vertex_interface* get_thread_reference_vertex(d1::wait_tree_vertex_interface* wc);
+TBB_EXPORT d1::task* __TBB_EXPORTED_FUNC current_task_ptr();
 
 // Do not place under __TBB_RESUMABLE_TASKS. It is a stub for unsupported platforms.
 struct suspend_point_type;
@@ -243,6 +244,7 @@ inline void wait(wait_context& wait_ctx, task_group_context& ctx) {
     call_itt_task_notify(destroy, &wait_ctx);
 }
 
+using r1::current_task_ptr;
 using r1::current_context;
 
 class task_traits {
diff --git a/src/tbb/def/lin32-tbb.def b/src/tbb/def/lin32-tbb.def
@@ -1,5 +1,6 @@
 /*
     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.
@@ -78,6 +79,7 @@ _ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE;
 _ZN3tbb6detail2r121current_suspend_pointEv;
 _ZN3tbb6detail2r114notify_waitersEj;
 _ZN3tbb6detail2r127get_thread_reference_vertexEPNS0_2d126wait_tree_vertex_interfaceE;
+_ZN3tbb6detail2r116current_task_ptrEv;
 
 /* Task dispatcher (task_dispatcher.cpp) */
 _ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE;
diff --git a/src/tbb/def/lin64-tbb.def b/src/tbb/def/lin64-tbb.def
@@ -1,5 +1,6 @@
 /*
     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.
@@ -78,6 +79,7 @@ _ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE;
 _ZN3tbb6detail2r121current_suspend_pointEv;
 _ZN3tbb6detail2r114notify_waitersEm;
 _ZN3tbb6detail2r127get_thread_reference_vertexEPNS0_2d126wait_tree_vertex_interfaceE;
+_ZN3tbb6detail2r116current_task_ptrEv;
 
 /* Task dispatcher (task_dispatcher.cpp) */
 _ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE;
diff --git a/src/tbb/def/mac64-tbb.def b/src/tbb/def/mac64-tbb.def
@@ -1,4 +1,5 @@
 # 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.
@@ -80,6 +81,7 @@ __ZN3tbb6detail2r16resumeEPNS1_18suspend_point_typeE
 __ZN3tbb6detail2r121current_suspend_pointEv
 __ZN3tbb6detail2r114notify_waitersEm
 __ZN3tbb6detail2r127get_thread_reference_vertexEPNS0_2d126wait_tree_vertex_interfaceE
+__ZN3tbb6detail2r116current_task_ptrEv
 
 # Task dispatcher (task_dispatcher.cpp)
 __ZN3tbb6detail2r114execution_slotEPKNS0_2d114execution_dataE
diff --git a/src/tbb/def/win32-tbb.def b/src/tbb/def/win32-tbb.def
@@ -1,4 +1,5 @@
 ; 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.
@@ -72,6 +73,7 @@ EXPORTS
 ?suspend@r1@detail@tbb@@YAXP6AXPAXPAUsuspend_point_type@123@@Z0@Z
 ?notify_waiters@r1@detail@tbb@@YAXI@Z
 ?get_thread_reference_vertex@r1@detail@tbb@@YAPAVwait_tree_vertex_interface@d1@23@PAV4523@@Z
+?current_task_ptr@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
diff --git a/src/tbb/def/win64-tbb.def b/src/tbb/def/win64-tbb.def
@@ -1,4 +1,5 @@
 ; 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.
@@ -72,6 +73,7 @@ EXPORTS
 ?current_suspend_point@r1@detail@tbb@@YAPEAUsuspend_point_type@123@XZ
 ?notify_waiters@r1@detail@tbb@@YAX_K@Z
 ?get_thread_reference_vertex@r1@detail@tbb@@YAPEAVwait_tree_vertex_interface@d1@23@PEAV4523@@Z
+?current_task_ptr@r1@detail@tbb@@YAPEAVtask@d1@23@XZ
 
 ; Task dispatcher (task_dispatcher.cpp)
 ?spawn@r1@detail@tbb@@YAXAEAVtask@d1@23@AEAVtask_group_context@523@@Z
diff --git a/src/tbb/scheduler_common.h b/src/tbb/scheduler_common.h
@@ -559,6 +559,8 @@ class alignas (max_nfs_size) task_dispatcher {
                       >
         m_reference_vertex_map;
 
+    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
         that has stolen it from the sender's task pool. Otherwise returns nullptr.
diff --git a/src/tbb/task.cpp b/src/tbb/task.cpp
@@ -256,6 +256,10 @@ d1::wait_tree_vertex_interface* get_thread_reference_vertex(d1::wait_tree_vertex
     return ref_counter;
 }
 
+d1::task* __TBB_EXPORTED_FUNC current_task_ptr() {
+    return governor::get_thread_data()->get_innermost_running_task();
+}
+
 } // namespace r1
 } // namespace detail
 } // namespace tbb
diff --git a/src/tbb/task_dispatcher.h b/src/tbb/task_dispatcher.h
@@ -254,11 +254,13 @@ 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;
         bool is_initially_registered;
 
         ~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;
 
             if (!is_initially_registered) {
                 task_disp.m_thread_data->my_arena->my_tc_client.get_pm_client()->unregister_thread();
@@ -268,7 +270,7 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
             __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, m_thread_data->my_is_registered };
+    } dl_guard{ *this, m_execute_data_ext, m_properties, m_innermost_running_task, m_thread_data->my_is_registered };
 
     // The context guard to track fp setting and itt tasks.
     context_guard_helper</*report_tasks=*/ITTPossible> context_guard;
@@ -331,6 +333,9 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
                     void* itt_caller = ed.context->my_itt_caller;
                     suppress_unused_warning(itt_caller);
 
+                    d1::task* prev_innermost_running_task = m_innermost_running_task;
+                    m_innermost_running_task = t;
+
                     ITT_CALLEE_ENTER(ITTPossible, t, itt_caller);
 
                     if (ed.context->is_group_execution_cancelled()) {
@@ -346,6 +351,7 @@ d1::task* task_dispatcher::local_wait_for_all(d1::task* t, Waiter& waiter ) {
                     ed.affinity_slot = d1::no_slot;
                     // Reset task owner id for bypassed task
                     ed.original_slot = m_thread_data->my_arena_index;
+                    m_innermost_running_task = prev_innermost_running_task;
                     t = get_critical_task(t, ed, isolation, critical_allowed);
                 }
                 __TBB_ASSERT(m_thread_data && governor::is_thread_data_set(m_thread_data), nullptr);
diff --git a/src/tbb/thread_data.h b/src/tbb/thread_data.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.
@@ -139,6 +140,7 @@ class thread_data : public ::rml::job
     void enter_task_dispatcher(task_dispatcher& task_disp, std::uintptr_t stealing_threshold);
     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);
+    d1::task* get_innermost_running_task();
 
     //! Index of the arena slot the scheduler occupies now, or occupied last time
     unsigned short my_arena_index;
@@ -254,9 +256,12 @@ 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 d1::task* thread_data::get_innermost_running_task() {
+    return my_task_dispatcher->m_innermost_running_task;
+}
+
 } // namespace r1
 } // namespace detail
 } // namespace tbb
 
 #endif // __TBB_thread_data_H
-
diff --git a/test/tbb/test_task.cpp b/test/tbb/test_task.cpp
@@ -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.
@@ -902,3 +903,141 @@ TEST_CASE("Force thread limit on per-thread reference_vertex") {
         groups[i].wait();
     }
 }
+
+constexpr std::size_t max_current_task_ptr_test_depth = 10;
+
+enum current_task_ptr_submit_type {
+    execute_and_wait,
+    submit_non_critical,
+    submit_critical,
+    slot_spawn,
+    enqueue
+};
+
+struct current_task_ptr_checking_task : public tbb::detail::d1::task {
+    current_task_ptr_checking_task(std::size_t d, tbb::task_arena& a, tbb::detail::d1::wait_context& wtc,
+                                   tbb::task_group_context& ctx, current_task_ptr_submit_type tag)
+        : depth(d), arena(a), wait_context(wtc), task_group_ctx(ctx), submit_tag(tag) {}
+
+    void submit_and_wait(current_task_ptr_checking_task& t, tbb::detail::d1::wait_context& wait_ctx) {
+        switch (submit_tag) {
+            case current_task_ptr_submit_type::execute_and_wait: {
+                tbb::detail::d1::execute_and_wait(t, task_group_ctx, wait_ctx, task_group_ctx);
+                break;
+            } case current_task_ptr_submit_type::submit_non_critical: {
+                submit(t, arena, task_group_ctx, false);
+                break;
+            } case current_task_ptr_submit_type::submit_critical: {
+                submit(t, arena, task_group_ctx, true);
+                break;
+            } case current_task_ptr_submit_type::slot_spawn: {
+                auto task_slot = tbb::detail::d1::slot_id(tbb::this_task_arena::current_thread_index() + 1 % tbb::this_task_arena::max_concurrency());
+                tbb::detail::d1::spawn(t, task_group_ctx, task_slot);
+                break; 
+            } case current_task_ptr_submit_type::enqueue: {
+                tbb::detail::r1::enqueue(t, task_group_ctx, &arena);
+                break;
+            } default:
+                CHECK(false);
+        }
+
+        if (submit_tag != current_task_ptr_submit_type::execute_and_wait) {
+            arena.execute([&] {
+                tbb::detail::r1::wait(wait_ctx, task_group_ctx);
+            });
+        }
+    }
+
+    tbb::detail::d1::task* execute(tbb::detail::d1::execution_data&) override {
+        ++execute_count;
+
+        REQUIRE_MESSAGE(tbb::detail::d1::current_task_ptr() == this,
+                        "Incorrect task returned from current_task_ptr");
+
+        if (depth < max_current_task_ptr_test_depth) {
+            tbb::detail::d1::wait_context next_task_wait_context(1);
+
+            current_task_ptr_checking_task next_task(depth + 1, arena, next_task_wait_context, task_group_ctx, submit_tag);
+            submit_and_wait(next_task, next_task_wait_context);
+
+            REQUIRE_MESSAGE(tbb::detail::d1::current_task_ptr() == this,
+                        "Incorrect task returned from current_task_ptr");
+        }
+
+        wait_context.release();
+        return nullptr;
+    }
+
+    tbb::detail::d1::task* cancel(tbb::detail::d1::execution_data&) override {
+        REQUIRE_MESSAGE(tbb::detail::d1::current_task_ptr() == this,
+                        "Incorrect task returned from current_task_ptr");
+        ++cancel_count;
+        wait_context.release();
+        return nullptr;
+    }
+
+    static std::atomic<std::size_t> execute_count;
+    static std::atomic<std::size_t> cancel_count;
+    std::size_t depth;
+    tbb::task_arena& arena;
+    tbb::detail::d1::wait_context& wait_context;
+    tbb::task_group_context& task_group_ctx;
+    current_task_ptr_submit_type submit_tag;
+};
+
+std::atomic<std::size_t> current_task_ptr_checking_task::execute_count{0};
+std::atomic<std::size_t> current_task_ptr_checking_task::cancel_count{0};
+
+void test_current_task_ptr_base(current_task_ptr_submit_type submit_tag, bool test_cancel) {
+    tbb::task_arena arena;
+    arena.initialize();
+    tbb::task_group_context test_context;
+
+    constexpr std::size_t num_start_tasks = 5;
+    tbb::detail::d1::wait_context root_task_wait(num_start_tasks);
+
+    REQUIRE_MESSAGE(tbb::detail::d1::current_task_ptr() == nullptr,
+                    "Incorrect task returned from current_task_ptr");
+
+    using task_type = current_task_ptr_checking_task;
+    std::vector<task_type, tbb::cache_aligned_allocator<task_type>> start_tasks;
+    start_tasks.reserve(num_start_tasks);
+
+    for (std::size_t i = 0; i < num_start_tasks; ++i) {
+        start_tasks.emplace_back(0, arena, root_task_wait, test_context, submit_tag);
+        submit(start_tasks.back(), arena, test_context, false);
+        if (test_cancel && i == num_start_tasks / 2) {
+            test_context.cancel_group_execution();
+        }
+    }
+
+    arena.execute([&] {
+        tbb::detail::d1::wait(root_task_wait, test_context);
+    });
+
+    std::size_t expected_tasks_count = max_current_task_ptr_test_depth * num_start_tasks + num_start_tasks;
+    if (test_cancel) {
+        REQUIRE_MESSAGE(task_type::cancel_count > 0, "Some tasks should be cancelled");
+        REQUIRE_MESSAGE(task_type::execute_count + task_type::cancel_count < expected_tasks_count,
+                        "Incorrect number of tasks executed or cancelled");
+    } else {
+        REQUIRE_MESSAGE(task_type::cancel_count == 0, "No tasks should be cancelled");
+        REQUIRE_MESSAGE(task_type::execute_count == expected_tasks_count, "Incorrect number of tasks executed");
+    }
+    task_type::execute_count = 0;
+    task_type::cancel_count = 0;
+}
+
+void test_current_task_ptr(bool test_cancel) {
+    test_current_task_ptr_base(current_task_ptr_submit_type::execute_and_wait, test_cancel);
+    test_current_task_ptr_base(current_task_ptr_submit_type::submit_non_critical, test_cancel);
+    test_current_task_ptr_base(current_task_ptr_submit_type::submit_critical, test_cancel);
+    test_current_task_ptr_base(current_task_ptr_submit_type::slot_spawn, test_cancel);
+    test_current_task_ptr_base(current_task_ptr_submit_type::enqueue, test_cancel);
+}
+
+//! \brief \ref error_guessing
+TEST_CASE("Test current_task_ptr") {
+    test_current_task_ptr(/*test_cancel = */false);
+    test_current_task_ptr(/*test_cancel = */true);
+}

Original file line number	Diff line number	Diff line change
`@@ -559,6 +559,8 @@ class alignas (max_nfs_size) task_dispatcher {`
`559`	`559`	`>`
`560`	`560`	`m_reference_vertex_map;`
`561`	`561`
	`562`	`+ d1::task* m_innermost_running_task{ nullptr };`
	`563`	`+`
`562`	`564`	`//! Attempt to get a task from the mailbox.`
`563`	`565`	`/** Gets a task only if it has not been executed by its sender or a thief`
`564`	`566`	`that has stolen it from the sender's task pool. Otherwise returns nullptr.`