[oneTBB] Add a function to create a set of NUMA bound task arenas (#650)

Signed-off-by: Isaev, Ilya <[email protected]>
Co-authored-by: Aleksei Fedotov <[email protected]>
Co-authored-by: Alexey Kukanov <[email protected]>
Co-authored-by: Konstantin Boyarinov <[email protected]>

Author: 4 people

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

@@ -42,18 +42,18 @@ A class that represents an explicit, user-managed task scheduler arena.
                     int max_threads_per_core = task_arena::automatic;
                 };
 
-                task_arena(int max_concurrency = automatic, unsigned reserved_for_masters = 1,
-                        priority a_priority = priority::normal);
-                task_arena(constraints a_constraints, unsigned reserved_for_masters = 1,
-                        priority a_priority = priority::normal);
+                task_arena(int max_concurrency = automatic, unsigned reserved_slots = 1,
+                           priority a_priority = priority::normal);
+                task_arena(constraints constraints_, unsigned reserved_slots = 1,
+                           priority a_priority = priority::normal);
                 task_arena(const task_arena &s);
                 explicit task_arena(oneapi::tbb::attach);
                 ~task_arena();
 
                 void initialize();
-                void initialize(int max_concurrency, unsigned reserved_for_masters = 1,
+                void initialize(int max_concurrency, unsigned reserved_slots = 1,
                                 priority a_priority = priority::normal);
-                void initialize(constraints a_constraints, unsigned reserved_for_masters = 1,
+                void initialize(constraints constraints_, unsigned reserved_slots = 1,
                                 priority a_priority = priority::normal);
                 void initialize(oneapi::tbb::attach);
 
@@ -71,6 +71,9 @@ A class that represents an explicit, user-managed task scheduler arena.
                 task_group_status task_arena::wait_for(task_group& tg);
             };
 
+            std::vector<task_arena> create_numa_task_arenas(task_arena::constraints constraints_ = {},
+                                                            unsigned reserved_slots = 0);
+
         } // namespace tbb
     } // namespace oneapi
 
@@ -182,28 +185,28 @@ Member types and constants
 Member functions
 ----------------
 
-.. cpp:function:: task_arena(int max_concurrency = automatic, unsigned reserved_for_masters = 1, priority a_priority = priority::normal)
+.. cpp:function:: task_arena(int max_concurrency = automatic, unsigned reserved_slots = 1, priority a_priority = priority::normal)
 
     Creates a ``task_arena`` with a certain concurrency limit (``max_concurrency``) and priority
     (``a_priority``).  Some portion of the limit can be reserved for application threads with
-    ``reserved_for_masters``.  The amount for reservation cannot exceed the limit.
+    ``reserved_slots``.  The amount for reservation cannot exceed the limit.
 
     .. caution::
 
-        If ``max_concurrency`` and ``reserved_for_masters`` are
+        If ``max_concurrency`` and ``reserved_slots`` are
         explicitly set to be equal and greater than 1, oneTBB worker threads will never
         join the arena. As a result, the execution guarantee for enqueued tasks is not valid
         in such arena. Do not use ``task_arena::enqueue()`` with an arena set to have no worker threads.
 
-.. cpp:function:: task_arena(constraints a_constraints, unsigned reserved_for_masters = 1, priority a_priority = priority::normal)
+.. cpp:function:: task_arena(constraints constraints_, unsigned reserved_slots = 1, priority a_priority = priority::normal)
 
-    Creates a ``task_arena`` with a certain constraints(``a_constraints``) and priority
+    Creates a ``task_arena`` with a certain constraints(``constraints_``) and priority
     (``a_priority``).  Some portion of the limit can be reserved for application threads with
-    ``reserved_for_masters``.  The amount for reservation cannot exceed the concurrency limit specified in ``constraints``.
+    ``reserved_slots``.  The amount for reservation cannot exceed the concurrency limit specified in ``constraints``.
 
     .. caution::
 
-        If ``constraints::max_concurrency`` and ``reserved_for_masters`` are
+        If ``constraints::max_concurrency`` and ``reserved_slots`` are
         explicitly set to be equal and greater than 1, oneTBB worker threads will never
         join the arena. As a result, the execution guarantee for enqueued tasks is not valid
         in such arena. Do not use ``task_arena::enqueue()`` with an arena set to have no worker threads.
@@ -239,11 +242,11 @@ Member functions
 
         After the call to ``initialize``, the arena parameters are fixed and cannot be changed.
 
-.. cpp:function:: void initialize(int max_concurrency, unsigned reserved_for_masters = 1, priority a_priority = priority::normal)
+.. cpp:function:: void initialize(int max_concurrency, unsigned reserved_slots = 1, priority a_priority = priority::normal)
 
     Same as above, but overrides previous arena parameters.
 
-.. cpp:function:: void initialize(constraints a_constraints, unsigned reserved_for_masters = 1, priority a_priority = priority::normal)
+.. cpp:function:: void initialize(constraints constraints_, unsigned reserved_slots = 1, priority a_priority = priority::normal)
 
     Same as above.
 
@@ -329,6 +332,26 @@ Member functions
 
     The behavior of this function is equivalent to ``this->execute([&tg]{ return tg.wait(); })``.
 
+Non-member Functions
+--------------------
+
+.. cpp:function:: std::vector<task_arena> create_numa_task_arenas(task_arena::constraints constraints_ = {}, unsigned reserved_slots = 0)
+
+    Returns a ``std::vector`` of non-initialized ``task_arena`` objects, each bound to a separate NUMA node.
+    The number of created ``task_arena`` instances is equal to the number of NUMA nodes on the system,
+    as determined by ``tbb::info::numa_nodes()``.
+    
+    If an error occurs during system information discovery,
+    returns a ``std::vector`` containing a single ``task_arena`` object created as 
+    ``task_arena(constraints_.set_numa_id(task_arena::automatic), reserved_slots)``.
+
+    The ``constraints_`` argument can be specified to apply additional limitations to threads in
+    the ``task_arena`` objects. For each created arena, the ``numa_id`` value in ``constraints_``
+    is automatically set to the corresponding NUMA node ID from ``tbb::info::numa_nodes()``.
+    
+    The ``reserved_slots`` argument allows reserving a specified number of slots in
+    each ``task_arena`` object for application threads. By default, no slots are reserved.
+
 Example
 -------
 
@@ -343,22 +366,21 @@ to the corresponding NUMA node.
     #include <vector>
 
     int main() {
-        std::vector<oneapi::tbb::numa_node_id> numa_nodes = oneapi::tbb::info::numa_nodes();
-        std::vector<oneapi::tbb::task_arena> arenas(numa_nodes.size());
-        std::vector<oneapi::tbb::task_group> task_groups(numa_nodes.size());
-
-        for (int i = 0; i < numa_nodes.size(); i++) {
-            arenas[i].initialize(oneapi::tbb::task_arena::constraints(numa_nodes[i]));
-        }
+        std::vector<oneapi::tbb::task_arena> arenas = oneapi::tbb::create_numa_task_arenas();
+        std::vector<oneapi::tbb::task_group> task_groups(arenas.size()-1);
 
-        for (int i = 0; i < numa_nodes.size(); i++) {
+        for (int i = 1; i < arenas.size(); i++) {
             arenas[i].enqueue([]{
                 /* executed by a thread pinned to the specified NUMA node */
-            }, task_groups[i]);
+            }, task_groups[i-1]);
         }
 
-        for (int i = 0; i < numa_nodes.size(); i++) {
-            arenas[i].wait_for(task_groups[i]);
+        arenas[0].execute([] {
+            /* executed by the main thread pinned to the NUMA node for arenas[0] */
+        });
+
+        for (int i = 1; i < arenas.size(); i++) {
+            arenas[i].wait_for(task_groups[i-1]);
         }
 
         return 0;