Implement multiple core type constraints using selectors (#1987)

Co-authored-by: Ilya Isaev <ilya.isaev@intel.com>

Author: dnmokhov

include/oneapi/tbb/detail/_config.h

@@ -553,11 +553,19 @@
 #define __TBB_PREVIEW_PARALLEL_PHASE 1
 #endif
 
+#if TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR || __TBB_BUILD
+#define __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR 1
+#endif
+
 #if !__TBB_DISABLE_SPEC_EXTENSIONS
 #define TBB_EXT_CUSTOM_ASSERTION_HANDLER 202510
 #endif
 
 // Feature-test macros
+#if __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR
+#define TBB_HAS_TASK_ARENA_CORE_TYPE_SELECTOR 202603
+#endif
+
 #if __TBB_PREVIEW_TASK_GROUP_EXTENSIONS
 #define TBB_HAS_TASK_GROUP_WAIT_FOR_SINGLE_TASK 202603
 #endif

include/oneapi/tbb/detail/_utils.h

@@ -22,6 +22,8 @@
 #include <cstdint>
 #include <atomic>
 #include <functional>
+#include <vector>
+#include <climits>
 
 #include "_config.h"
 #include "_assert.h"
@@ -156,6 +158,59 @@ T reverse_n_bits(T src, std::size_t n) {
     return reverse_bits(src) >> (number_of_bits<T>() - n);
 }
 
+//! Encodes/decodes multiple core type IDs into/from a single integer value using bitmask
+struct multi_core_type_codec {
+    using core_type_id = int;
+    static constexpr core_type_id automatic = -1;
+
+    static core_type_id encode(const std::vector<core_type_id>& ids) {
+        if (ids.empty()) {
+            return automatic;
+        }
+        if (ids.size() == 1) {
+            return ids[0];
+        }
+
+        core_type_id result = (encoding_format << bitmask_width);
+
+        for (core_type_id id : ids) {
+            __TBB_ASSERT((0 <= id) && (id < static_cast<core_type_id>(bitmask_width)), "Wrong core type id");
+            result |= (1 << id);
+        }
+
+        return result;
+    }
+    static std::vector<core_type_id> decode(core_type_id core_type) {
+        if (!is_encoded(core_type)) {
+            return {core_type};
+        }
+
+        std::vector<core_type_id> core_type_ids;
+        for (size_t bit_pos = 0; bit_pos < bitmask_width; ++bit_pos) {
+            if (core_type & (1 << bit_pos)) {
+                core_type_ids.push_back(static_cast<core_type_id>(bit_pos));
+            }
+        }
+        return core_type_ids;
+    }
+    static bool is_single(core_type_id id) {
+        return (id >> bitmask_width) == 0;
+    }
+    static bool is_encoded(core_type_id id) {
+        return (static_cast<std::make_unsigned<core_type_id>::type>(id) >> bitmask_width) == encoding_format;
+    }
+    static bool is_core_type(core_type_id id) {
+        return is_single(id) || is_encoded(id);
+    }
+
+    // Lower bitmask_width bits encode IDs
+    static constexpr size_t bitmask_width = sizeof(core_type_id) * CHAR_BIT - 4;
+
+    // Upper 4 bits: MSb=1 (makes result negative; real core type IDs are non-negative) + 3-bit format version
+    // (current: 0, max: 6; 1111 is excluded to avoid collision with plain negatives: -1, -2, ..., -268435456)
+    static constexpr size_t encoding_format = 0x8;
+};
+
 // A function to check if passed integer is a power of two
 template <typename IntegerType>
 constexpr bool is_power_of_two( IntegerType arg ) {

include/oneapi/tbb/info.h

@@ -1,5 +1,6 @@
 /*
     Copyright (c) 2019-2022 Intel Corporation
+    Copyright (c) 2026 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.
@@ -19,8 +20,11 @@
 
 #include "detail/_config.h"
 #include "detail/_namespace_injection.h"
+#include "detail/_utils.h"
+#include "version.h"
 
 #if __TBB_ARENA_BINDING
+#include <tuple>
 #include <vector>
 #include <cstdint>
 
@@ -103,6 +107,59 @@ inline int default_concurrency(constraints c) {
     return r1::constraints_default_concurrency(c);
 }
 
+#if __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR
+// Call a custom selector on the available core type(s) and encode those selected
+template <typename Selector>
+inline core_type_id apply_core_type_selector(Selector selector) {
+    constexpr core_type_id automatic = -1;
+
+    auto ids = core_types();
+    size_t total = ids.size();
+    if (total < 2) {
+        // Not enough core types to select from, so use the default
+        return automatic;
+    }
+
+    int max_score = 0, max_score_id = -1, num_zero_scores = 0;
+    std::vector<core_type_id> selected_core_types;
+    for (size_t index = 0; index < total; ++index) {
+        int score = selector(std::make_tuple(ids[index], index, total));
+        if (score > 0) {
+            selected_core_types.push_back(ids[index]);
+        }
+        else if (score == 0) {
+            ++num_zero_scores;
+        }
+
+        if (TBB_runtime_interface_version() < 12180) {
+            if (score > max_score) {
+                max_score = score;
+                max_score_id = ids[index];
+            }
+        }
+    }
+    if (TBB_runtime_interface_version() < 12180) {
+        // No runtime multi core type support, so select all or one
+        if (selected_core_types.size() + num_zero_scores == total) {
+            selected_core_types.clear(); // all
+        }
+        else if (!selected_core_types.empty()) {
+            selected_core_types = { max_score_id }; // the one with the highest score
+        }
+    }
+    return multi_core_type_codec::encode(selected_core_types);
+}
+
+template <typename Selector>
+inline int default_concurrency(constraints c, Selector selector) {
+    constexpr core_type_id selectable = -2;
+    if (c.core_type == selectable) {
+        c.core_type = apply_core_type_selector(selector);
+    }
+    return default_concurrency(c);
+}
+#endif
+
 } // namespace d1
 } // namespace detail
 

include/oneapi/tbb/task_arena.h

@@ -253,7 +253,12 @@ class task_arena_base {
 public:
     //! Typedef for number of threads that is automatic.
     static const int automatic = -1;
+    //! Typedef for current thread index in an uninitialized arena.
     static const int not_initialized = -2;
+#if __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR
+    //! Typedef for core type(s) to be specified by the provided selector.
+    static const int selectable = -2;
+#endif
 };
 
 template<typename R, typename F>
@@ -364,6 +369,28 @@ class task_arena : public task_arena_base {
           )
     {}
 
+#if __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR
+    //! Creates task arena with a custom selector for core types
+    template <typename Selector,
+              typename = decltype(static_cast<int>(std::declval<Selector>()(std::declval<std::tuple<int, size_t, size_t>>())))>
+    task_arena(const constraints& constraints_, Selector selector_,
+               unsigned reserved_for_masters = 1, priority a_priority = priority::normal
+#if __TBB_PREVIEW_PARALLEL_PHASE
+               , leave_policy lp = leave_policy::automatic
+#endif
+    )
+        : task_arena_base(constraints_, reserved_for_masters, a_priority
+#if __TBB_PREVIEW_PARALLEL_PHASE
+                         , lp
+#endif
+          )
+    {
+        if (my_core_type == selectable) {
+            my_core_type = apply_core_type_selector(selector_);
+        }
+    }
+#endif
+
     //! Copies settings from another task_arena
     task_arena(const task_arena &a) // copy settings but not the reference or instance
         : task_arena_base(
@@ -440,6 +467,7 @@ class task_arena : public task_arena_base {
     }
 
 #if __TBB_ARENA_BINDING
+    //! Overrides constraints and forces initialization of internal representation
     void initialize(constraints constraints_, unsigned reserved_slots = 1,
                     priority a_priority = priority::normal
 #if __TBB_PREVIEW_PARALLEL_PHASE
@@ -462,6 +490,37 @@ class task_arena : public task_arena_base {
             mark_initialized();
         }
     }
+
+#if __TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR
+    //! Overrides constraints with a custom selector for core types and forces initialization of internal representation
+    template<typename Selector,
+             typename = decltype(static_cast<int>(std::declval<Selector>()(std::declval<std::tuple<int, size_t, size_t>>())))>
+    void initialize(constraints constraints_, Selector selector_,
+                    unsigned reserved_for_masters = 1, priority a_priority = priority::normal
+#if __TBB_PREVIEW_PARALLEL_PHASE
+                    , leave_policy lp = leave_policy::automatic
+#endif
+    )
+    {
+        __TBB_ASSERT(!my_arena.load(std::memory_order_relaxed), "Impossible to modify settings of an already initialized task_arena");
+        if( !is_active() ) {
+            my_numa_id = constraints_.numa_id;
+            my_max_concurrency = constraints_.max_concurrency;
+            my_core_type = constraints_.core_type;
+            my_max_threads_per_core = constraints_.max_threads_per_core;
+            my_num_reserved_slots = reserved_for_masters;
+            my_priority = a_priority;
+#if __TBB_PREVIEW_PARALLEL_PHASE
+            set_leave_policy(lp);
+#endif
+            if (my_core_type == selectable) {
+                my_core_type = apply_core_type_selector(selector_);
+            }
+            r1::initialize(*this);
+            mark_initialized();
+        }
+    }
+#endif /*__TBB_PREVIEW_TASK_ARENA_CORE_TYPE_SELECTOR*/
 #endif /*__TBB_ARENA_BINDING*/
 
     //! Attaches this instance to the current arena of the thread

src/tbb/arena.cpp

@@ -58,7 +58,7 @@ class numa_binding_observer : public tbb::task_scheduler_observer {
 
 numa_binding_observer* construct_binding_observer( d1::task_arena* ta, int num_slots, int numa_id, core_type_id core_type, int max_threads_per_core ) {
     numa_binding_observer* binding_observer = nullptr;
-    if ((core_type >= 0 && core_type_count() > 1) || (numa_id >= 0 && numa_node_count() > 1) || max_threads_per_core > 0) {
+    if ((multi_core_type_codec::is_core_type(core_type) && core_type_count() > 1) || (numa_id >= 0 && numa_node_count() > 1) || max_threads_per_core > 0) {
         binding_observer = new(allocate_memory(sizeof(numa_binding_observer))) numa_binding_observer(ta, num_slots, numa_id, core_type, max_threads_per_core);
         __TBB_ASSERT(binding_observer, "Failure during NUMA binding observer allocation and construction");
     }

src/tbb/governor.cpp

@@ -1,6 +1,6 @@
 /*
     Copyright (c) 2005-2025 Intel Corporation
-    Copyright (c) 2025 UXL Foundation Contributors
+    Copyright (c) 2025-2026 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.
@@ -591,7 +591,7 @@ void constraints_assertion(d1::constraints c) {
     int* core_types_begin = system_topology::core_types_indexes;
     int* core_types_end = system_topology::core_types_indexes + system_topology::core_types_count;
     __TBB_ASSERT_RELEASE(c.core_type == system_topology::automatic ||
-        (is_topology_initialized && std::find(core_types_begin, core_types_end, c.core_type) != core_types_end),
+        (is_topology_initialized && (multi_core_type_codec::is_encoded(c.core_type) || std::find(core_types_begin, core_types_end, c.core_type) != core_types_end)),
         "The constraints::core_type value is not known to the library. Use tbb::info::core_types() to get the list of possible values.");
 }
 
@@ -600,7 +600,7 @@ int __TBB_EXPORTED_FUNC constraints_default_concurrency(const d1::constraints& c
 
     const int default_num_threads = int(governor::default_num_threads());
 
-    if (c.numa_id >= 0 || c.core_type >= 0 || c.max_threads_per_core > 0) {
+    if (c.numa_id >= 0 || multi_core_type_codec::is_core_type(c.core_type) || c.max_threads_per_core > 0) {
         system_topology::initialize();
         const int constrained_default_concurrency =
             get_default_concurrency_ptr(c.numa_id, c.core_type, c.max_threads_per_core);

src/tbbbind/tbb_bind.cpp

@@ -1,6 +1,6 @@
 /*
     Copyright (c) 2019-2025 Intel Corporation
-    Copyright (c) 2025 UXL Foundation Contributors
+    Copyright (c) 2025-2026 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.
@@ -21,6 +21,7 @@
 #include "../tbb/assert_impl.h" // Out-of-line TBB assertion handling routines are instantiated here.
 #include "oneapi/tbb/detail/_assert.h"
 #include "oneapi/tbb/detail/_config.h"
+#include "oneapi/tbb/detail/_utils.h"
 
 #if _MSC_VER && !__INTEL_COMPILER && !__clang__
 #pragma warning( push )
@@ -355,7 +356,13 @@ class system_topology {
     void fill_constraints_affinity_mask(affinity_mask input_mask, int numa_node_index, int core_type_index, int max_threads_per_core) {
         __TBB_ASSERT(is_topology_parsed(), "Trying to get access to uninitialized system_topology");
         __TBB_ASSERT(numa_node_index < (int)numa_affinity_masks_list.size(), "Wrong NUMA node id");
-        __TBB_ASSERT(core_type_index < (int)core_types_affinity_masks_list.size(), "Wrong core type id");
+        __TBB_ASSERT(core_type_index == -1 ||
+            // In the multiple core type format, the MSb of the first bitmask_width bits represents the highest core type id
+            (multi_core_type_codec::is_single(core_type_index)
+                 ? (size_t)core_type_index
+                 : log2(core_type_index & ((1 << multi_core_type_codec::bitmask_width) - 1))) <
+                core_types_affinity_masks_list.size(),
+            "Wrong core type id");
         __TBB_ASSERT(max_threads_per_core == -1 || max_threads_per_core > 0, "Wrong max_threads_per_core");
 
         hwloc_cpuset_t constraints_mask = hwloc_bitmap_alloc();
@@ -365,8 +372,19 @@ class system_topology {
         if (numa_node_index >= 0) {
             hwloc_bitmap_and(constraints_mask, constraints_mask, numa_affinity_masks_list[numa_node_index]);
         }
-        if (core_type_index >= 0) {
-            hwloc_bitmap_and(constraints_mask, constraints_mask, core_types_affinity_masks_list[core_type_index]);
+        if (multi_core_type_codec::is_core_type(core_type_index)) {
+            auto core_types = multi_core_type_codec::decode(core_type_index);
+            __TBB_ASSERT(!core_types.empty(), "Core types list must not be empty");
+
+            hwloc_cpuset_t core_types_mask = hwloc_bitmap_alloc();
+
+            // Combine affinity masks for specified core types
+            for (int c : core_types) {
+                hwloc_bitmap_or(core_types_mask, core_types_mask, core_types_affinity_masks_list[c]);
+            }
+
+            hwloc_bitmap_and(constraints_mask, constraints_mask, core_types_mask);
+            hwloc_bitmap_free(core_types_mask);
         }
         if (max_threads_per_core > 0) {
             // clear input mask