[oneTBB] Improve interoperability of task arenas and task groups (#624)

Author: akukanov

source/elements/oneTBB/source/task_scheduler/task_arena/task_arena_cls.rst

@@ -1,4 +1,5 @@
 .. SPDX-FileCopyrightText: 2019-2021 Intel Corporation
+.. SPDX-FileCopyrightText: 2025 UXL Foundation Contributors
 ..
 .. SPDX-License-Identifier: CC-BY-4.0
 
@@ -62,9 +63,12 @@ A class that represents an explicit, user-managed task scheduler arena.
                 int max_concurrency() const;
 
                 template<typename F> auto execute(F&& f) -> decltype(f());
-                template<typename F> void enqueue(F&& f);
 
+                template<typename F> void enqueue(F&& f);
+                template<typename F> void enqueue(F&& f, task_group& tg);
                 void enqueue(task_handle&& h);
+
+                task_group_status task_arena::wait_for(task_group& tg);
             };
 
         } // namespace tbb
@@ -263,7 +267,26 @@ Member functions
     Returns the concurrency level of the ``task_arena``.
     Does not require the ``task_arena`` to be initialized and does not perform initialization.
 
-.. cpp:function:: template<F> void enqueue(F&& f)
+.. cpp:function:: template<typename F> auto execute(F&& f) -> decltype(f())
+
+    Executes the specified functor in the ``task_arena`` and returns the value returned by the functor.
+    The ``F`` type must meet the `Function Objects` requirements described in the [function.objects] section of the ISO C++ standard.
+
+    The calling thread joins the ``task_arena`` if possible, and executes the functor.
+    Upon return it restores the previous task scheduler state and floating-point settings.
+
+    If joining the ``task_arena`` is not possible, the call wraps the functor into a task,
+    enqueues it into the arena, waits using an OS kernel synchronization object
+    for another opportunity to join, and finishes after the task completion.
+
+    An exception thrown in the functor will be captured and re-thrown from ``execute``.
+
+    .. note::
+
+        Any number of threads outside of the arena can submit work to the arena and be blocked.
+        However, only the maximal number of threads specified for the arena can participate in executing the work.
+
+.. cpp:function:: template<typename F> void enqueue(F&& f)
 
     Enqueues a task into the ``task_arena`` to process the specified functor and immediately returns.
     The ``F`` type must meet the `Function Objects` requirements described in the [function.objects] section of the ISO C++ standard.
@@ -284,34 +307,28 @@ Member functions
 
         An exception thrown and not caught in the functor results in undefined behavior.
 
-.. cpp:function:: template<F> auto execute(F&& f) -> decltype(f())
+.. cpp:function:: template<typename F> void enqueue(F&& f, task_group& tg)
 
-    Executes the specified functor in the ``task_arena`` and returns the value returned by the functor.
-    The ``F`` type must meet the `Function Objects` requirements described in the [function.objects] section of the ISO C++ standard.
-
-    The calling thread joins the ``task_arena`` if possible, and executes the functor.
-    Upon return it restores the previous task scheduler state and floating-point settings.
+    Adds a task to process the specified functor into ``tg`` and enqueues it into the ``task_arena``.
 
-    If joining the ``task_arena`` is not possible, the call wraps the functor into a task,
-    enqueues it into the arena, waits using an OS kernel synchronization object
-    for another opportunity to join, and finishes after the task completion.
-
-    An exception thrown in the functor will be captured and re-thrown from ``execute``.
-
-    .. note::
-
-        Any number of threads outside of the arena can submit work to the arena and be blocked.
-        However, only the maximal number of threads specified for the arena can participate in executing the work.
+    The behavior of this function is equivalent to ``this->enqueue( tg.defer(std::forward<F>(f)) )``.
 
 .. cpp:function:: void enqueue(task_handle&& h)   
      
     Enqueues a task owned by ``h`` into the ``task_arena`` for processing. 
 
-    The behavior of this function is identical to the generic version (``template<typename F> void task_arena::enqueue(F&& f)``), except parameter type. 
+    The behavior of this function is equivalent to the generic version (``template<typename F> void task_arena::enqueue(F&& f)``), except parameter type. 
 
     .. note:: 
        ``h`` should not be empty to avoid an undefined behavior.
 
+.. cpp:function:: task_group_status task_arena::wait_for(task_group& tg)
+
+    Waits for all tasks in ``tg`` to complete or be cancelled, while possibly executing tasks in the ``task_arena``.
+    Returns the status of ``tg`` once waiting is complete.
+
+    The behavior of this function is equivalent to ``this->execute([&tg]{ return tg.wait(); })``.
+
 Example
 -------
 
@@ -335,17 +352,13 @@ to the corresponding NUMA node.
         }
 
         for (int i = 0; i < numa_nodes.size(); i++) {
-            arenas[i].execute([&task_groups, i] {
-                task_groups[i].run([] {
-                    /* executed by the thread pinned to specified NUMA node */
-                });
-            });
+            arenas[i].enqueue([]{
+                /* executed by a thread pinned to the specified NUMA node */
+            }, task_groups[i]);
         }
 
         for (int i = 0; i < numa_nodes.size(); i++) {
-            arenas[i].execute([&task_groups, i] {
-                task_groups[i].wait();
-            });
+            arenas[i].wait_for(task_groups[i]);
         }
 
         return 0;

source/elements/oneTBB/source/task_scheduler/task_arena/this_task_arena_ns.rst

@@ -1,4 +1,5 @@
 .. SPDX-FileCopyrightText: 2019-2021 Intel Corporation
+.. SPDX-FileCopyrightText: 2025 UXL Foundation Contributors
 ..
 .. SPDX-License-Identifier: CC-BY-4.0
 
@@ -21,11 +22,12 @@ with the ``task_arena`` currently used by the calling thread.
         namespace this_task_arena {
             int current_thread_index();
             int max_concurrency();
+
             template<typename F> auto isolate(F&& f) -> decltype(f());
             
+            template<typename F> void enqueue(F&& f);
+            template<typename F> void enqueue(F&& f, task_group& tg);
             void enqueue(task_handle&& h);
-            
-            template<typename F> void enqueue(F&& f) ;
         }
     } // namespace tbb
     } // namespace oneapi 
@@ -63,7 +65,7 @@ with the ``task_arena`` currently used by the calling thread.
     If the thread has not yet initialized the task scheduler, returns the concurrency level
     determined automatically for the hardware configuration.
 
-.. cpp:function:: template<F> auto isolate(F&& f) -> decltype(f())
+.. cpp:function:: template<typename F> auto isolate(F&& f) -> decltype(f())
 
     Runs the specified functor in isolation by restricting the calling thread to process only tasks
     scheduled in the scope of the functor (also called the isolation region). The function returns the value returned by the functor.
@@ -77,15 +79,21 @@ with the ``task_arena`` currently used by the calling thread.
   
     Enqueues a task into the ``task_arena`` currently used by the calling thread to process the specified functor, then returns immediately.
     The ``F`` type must meet the `Function Objects` requirements described in the [function.objects] section of the ISO C++ standard.
-    
-    Behavior of this function is identical to ``template<typename F> void task_arena::enqueue(F&& f)`` applied to the ``task_arena`` 
-    object constructed with ``attach`` parameter.     
+
+    Behavior of this function is equivalent to ``template<typename F> void task_arena::enqueue(F&& f)`` applied to the ``task_arena`` 
+    object constructed with ``attach`` parameter.
+
+.. cpp:function:: template<typename F> void enqueue(F&& f, task_group& tg)
+  
+    Adds a task to process the specified functor into ``tg`` and enqueues it into the ``task_arena`` currently used by the calling thread.
+
+    The behavior of this function is equivalent to ``this_task_arena::enqueue( tg.defer(std::forward<F>(f)) )``.
 
 .. cpp:function:: void enqueue(task_handle&& h)   
      
     Enqueues a task owned by ``h`` into the ``task_arena`` that is currently used by the calling thread.
 
-    The behavior of this function is identical to the generic version (``template<typename F> void enqueue(F&& f)``), except the parameter type. 
+    The behavior of this function is equivalent to the generic version (``template<typename F> void enqueue(F&& f)``), except the parameter type. 
 
     .. note:: 
         ``h`` should not be empty to avoid an undefined behavior.