Changing threading control interfaces (WIP)

cpp/daal/BUILD

@@ -17,7 +17,7 @@ daal_module(
                                     "@micromkl//:vml_ipp",
                                     # TODO: Currently vml_ipp lib depends on TBB, but it shouldn't
                                     #       Remove TBB from deps once problem with vml_ipp is resolved
-                                    "@tbb//:tbb_binary",
+                                    "@tbb//:tbb",
                                 ],
         }),
 )

cpp/daal/src/services/env_detect.cpp

@@ -124,22 +124,10 @@ DAAL_EXPORT daal::services::Environment::Environment() : _globalControl {}
 
 DAAL_EXPORT daal::services::Environment::Environment(const Environment & e) : daal::services::Environment::Environment() {}
 
+/*This functionality has been moved to oneDAL side*/
 DAAL_EXPORT void daal::services::Environment::initNumberOfThreads()
 {
     if (isInit) return;
-
-    /* if HT enabled - set _numThreads to physical cores num */
-    if (daal::internal::ServiceInst::serv_get_ht())
-    {
-        /* Number of cores = number of cpu packages * number of cores per cpu package */
-        int ncores = daal::internal::ServiceInst::serv_get_ncpus() * daal::internal::ServiceInst::serv_get_ncorespercpu();
-
-        /*  Re-set number of threads if ncores is valid and different to _numThreads */
-        if ((ncores > 0) && (ncores < _daal_threader_get_max_threads()))
-        {
-            daal::services::Environment::setNumberOfThreads(ncores);
-        }
-    }
     isInit = true;
 }
 

cpp/daal/src/threading/service_thread_pinner.cpp

@@ -209,6 +209,7 @@ class thread_pinner_impl_t : public tbb::task_scheduler_observer
 
 public:
     thread_pinner_impl_t(void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *));
+    thread_pinner_impl_t(void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *), tbb::task_arena & task_arena);
     void on_scheduler_entry(bool);
     void on_scheduler_exit(bool);
     void init_thread_pinner(int statusToSet, int nthreadsToSet, int max_threadsToSet, int * cpu_queueToSet);
@@ -230,6 +231,7 @@ class thread_pinner_impl_t : public tbb::task_scheduler_observer
     int get_status();
     bool get_pinning();
     bool set_pinning(bool p);
+    tbb::task_arena * get_task_arena();
     ~thread_pinner_impl_t();
 } * IMPL;
 
@@ -245,6 +247,18 @@ thread_pinner_impl_t::thread_pinner_impl_t(void (*read_topo)(int &, int &, int &
     return;
 } /* thread_pinner_impl_t() */
 
+thread_pinner_impl_t::thread_pinner_impl_t(void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *), tbb::task_arena & task_arena)
+    : pinner_arena(task_arena), tbb::task_scheduler_observer(pinner_arena), topo_deleter(deleter)
+{
+    do_pinning = (nthreads > 0) ? true : false;
+    is_pinning.set(0);
+
+    read_topo(status, nthreads, max_threads, &cpu_queue);
+    observe(true);
+
+    return;
+} /* thread_pinner_impl_t() */
+
 void thread_pinner_impl_t::on_scheduler_entry(bool) /*override*/
 {
     if (do_pinning == false || status < 0) return;
@@ -324,6 +338,11 @@ bool thread_pinner_impl_t::set_pinning(bool p)
     return old_pinning;
 } /* bool set_pinning(bool p) */
 
+tbb::task_arena * thread_pinner_impl_t::get_task_arena()
+{
+    return &(this->pinner_arena);
+}
+
 thread_pinner_impl_t::~thread_pinner_impl_t()
 {
     observe(false);
@@ -352,12 +371,37 @@ DAAL_EXPORT void * _getThreadPinner(bool create_pinner, void (*read_topo)(int &,
     return NULL;
 } /* thread_pinner_t* getThreadPinner() */
 
+DAAL_EXPORT void * _getThreadPinnerFromTaskArena(bool create_pinner, void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *),
+                                                 tbb::task_arena & task_arena)
+{
+    static bool pinner_created = false;
+
+    if (create_pinner == true || pinner_created == true)
+    {
+        static daal::services::internal::thread_pinner_t thread_pinner(read_topo, deleter, task_arena);
+        if (thread_pinner.get_status() == 0)
+        {
+            pinner_created = true;
+            return (void *)&thread_pinner;
+        }
+    }
+
+    return NULL;
+} /* thread_pinner_t* getThreadPinner() */
+
 DAAL_EXPORT void _thread_pinner_thread_pinner_init(void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *))
 {
     static thread_pinner_impl_t impl(read_topo, deleter);
     IMPL = &impl;
 }
 
+DAAL_EXPORT void _thread_pinner_init_with_task_arena(void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *),
+                                                     tbb::task_arena & task_arena)
+{
+    static thread_pinner_impl_t impl(read_topo, deleter, task_arena);
+    IMPL = &impl;
+}
+
 DAAL_EXPORT void _thread_pinner_execute(daal::services::internal::thread_pinner_task_t & task)
 {
     IMPL->execute(task);
@@ -388,6 +432,11 @@ DAAL_EXPORT void _thread_pinner_on_scheduler_exit(bool p)
     IMPL->on_scheduler_exit(p);
 }
 
+DAAL_EXPORT tbb::task_arena * _thread_pinner_get_task_arena()
+{
+    return IMPL->get_task_arena();
+}
+
     #else /* if __DO_TBB_LAYER__ is not defined */
 
 DAAL_EXPORT void * _getThreadPinner(bool create_pinner, void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *))

cpp/daal/src/threading/service_thread_pinner.h

@@ -25,6 +25,7 @@
 #define __SERVICE_THREAD_PINNER_H__
 
 #include "services/daal_defines.h"
+#include "tbb/tbb.h"
 #if !defined(DAAL_THREAD_PINNING_DISABLED)
 
     #include <cstdlib>
@@ -89,6 +90,7 @@ class thread_pinner_task_t
 extern "C"
 {
     DAAL_EXPORT void _thread_pinner_thread_pinner_init(void(int &, int &, int &, int **), void (*deleter)(void *));
+    DAAL_EXPORT void _thread_pinner_init_with_task_arena(void(int &, int &, int &, int **), void (*deleter)(void *), tbb::task_arena & task_arena);
     DAAL_EXPORT void _thread_pinner_read_topology();
     DAAL_EXPORT void _thread_pinner_on_scheduler_entry(bool);
     DAAL_EXPORT void _thread_pinner_on_scheduler_exit(bool);
@@ -98,7 +100,11 @@ extern "C"
     DAAL_EXPORT bool _thread_pinner_get_pinning();
     DAAL_EXPORT bool _thread_pinner_set_pinning(bool p);
 
+    DAAL_EXPORT tbb::task_arena * _thread_pinner_get_task_arena();
+
     DAAL_EXPORT void * _getThreadPinner(bool create_pinner, void(int &, int &, int &, int **), void (*deleter)(void *));
+    DAAL_EXPORT void * _getThreadPinnerFromTaskArena(bool create_pinner, void(int &, int &, int &, int **), void (*deleter)(void *),
+                                                     tbb::task_arena & task_arena);
 }
 
 namespace daal
@@ -111,20 +117,33 @@ class thread_pinner_t
 {
 public:
     thread_pinner_t(void (*f)(int &, int &, int &, int **), void (*deleter)(void *)) { _thread_pinner_thread_pinner_init(f, deleter); }
+    thread_pinner_t(void (*f)(int &, int &, int &, int **), void (*deleter)(void *), tbb::task_arena & task_arena)
+    {
+        _thread_pinner_init_with_task_arena(f, deleter, task_arena);
+    }
     void read_topology() { _thread_pinner_read_topology(); }
     void on_scheduler_entry(bool p) { _thread_pinner_on_scheduler_entry(p); }
     void on_scheduler_exit(bool p) { _thread_pinner_on_scheduler_exit(p); }
     void execute(thread_pinner_task_t & task) { _thread_pinner_execute(task); }
     int get_status() { return _thread_pinner_get_status(); }
     bool get_pinning() { return _thread_pinner_get_pinning(); }
     bool set_pinning(bool p) { return _thread_pinner_set_pinning(p); }
+    tbb::task_arena * get_task_arena() { return _thread_pinner_get_task_arena(); }
 };
 
 inline thread_pinner_t * getThreadPinner(bool create_pinner, void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *))
 {
     return (thread_pinner_t *)_getThreadPinner(create_pinner, read_topo, deleter);
 }
 
+inline thread_pinner_t * getThreadPinnerFromTaskArena(bool create_pinner, void (*read_topo)(int &, int &, int &, int **), void (*deleter)(void *),
+                                                      tbb::task_arena & task_arena)
+{
+    return (thread_pinner_t *)_getThreadPinnerFromTaskArena(create_pinner, read_topo, deleter, task_arena);
+}
+
+class thread_pinner_impl_t;
+
 } // namespace internal
 } // namespace services
 } // namespace daal

cpp/oneapi/dal/BUILD

@@ -28,6 +28,7 @@ dal_module(
         "@onedal//cpp/oneapi:include_root",
         "@onedal//cpp/daal:services",
         "@onedal//cpp/daal:data_management",
+        "@onedal//cpp/daal:threading_tbb",
     ],
     dpc_deps = [
         "@micromkl_dpc//:mkl_dpc",

cpp/oneapi/dal/backend/dispatcher.hpp

@@ -22,6 +22,7 @@
 #include "oneapi/dal/backend/common.hpp"
 #include "oneapi/dal/backend/communicator.hpp"
 #include "oneapi/dal/backend/dispatcher_cpu.hpp"
+#include "oneapi/dal/backend/threading.hpp"
 
 #define KERNEL_SPEC(spec, ...) ::oneapi::dal::backend::kernel_spec<spec, __VA_ARGS__>
 
@@ -72,27 +73,43 @@ class communicator_provider : public base {
 class context_cpu : public communicator_provider<spmd::device_memory_access::none> {
 public:
     explicit context_cpu(const detail::host_policy& policy = detail::host_policy::get_default())
-            : cpu_extensions_(policy.get_enabled_cpu_extensions()) {
+            : cpu_extensions_(policy.get_enabled_cpu_extensions()),
+              threading_policy_(policy.get_threading_policy()) {
         global_init();
     }
 
     explicit context_cpu(const detail::spmd_host_policy& policy)
             : communicator_provider<spmd::device_memory_access::none>(policy.get_communicator()),
-              cpu_extensions_(policy.get_local().get_enabled_cpu_extensions()) {
+              cpu_extensions_(policy.get_local().get_enabled_cpu_extensions()),
+              threading_policy_(policy.get_local().get_threading_policy()) {
         global_init();
     }
 
     explicit context_cpu(const spmd::communicator<spmd::device_memory_access::none>& comm)
             : communicator_provider<spmd::device_memory_access::none>(comm),
-              cpu_extensions_(detail::host_policy::get_default().get_enabled_cpu_extensions()) {}
+              cpu_extensions_(detail::host_policy::get_default().get_enabled_cpu_extensions()),
+              threading_policy_(detail::host_policy::get_default().get_threading_policy()) {}
 
     detail::cpu_extension get_enabled_cpu_extensions() const {
         return cpu_extensions_;
     }
 
+    bool get_thread_pinning() const {
+        return threading_policy_.thread_pinning;
+    }
+
+    int get_max_threads_per_core() const {
+        return threading_policy_.max_threads_per_core;
+    }
+
+    threading_policy get_threading_policy() const {
+        return threading_policy_;
+    }
+
 private:
     void global_init();
     detail::cpu_extension cpu_extensions_;
+    detail::threading_policy threading_policy_;
 };
 
 #ifdef ONEDAL_DATA_PARALLEL
@@ -275,18 +292,32 @@ inline bool test_cpu_extension(detail::cpu_extension mask, detail::cpu_extension
 }
 
 template <typename Op>
-inline constexpr auto dispatch_by_cpu(const context_cpu& ctx, Op&& op) {
+inline auto dispatch_by_cpu(const context_cpu& ctx, Op&& op) {
     using detail::cpu_extension;
+    using detail::threading_policy;
+
+    threading_policy policy = ctx.get_threading_policy();
+    task_executor task_executor_(policy);
 
     [[maybe_unused]] const cpu_extension cpu_ex = ctx.get_enabled_cpu_extensions();
     ONEDAL_IF_CPU_DISPATCH_AVX512(if (test_cpu_extension(cpu_ex, cpu_extension::avx512)) {
-        return op(cpu_dispatch_avx512{});
+        return task_executor_.execute([&]() {
+            return op(cpu_dispatch_avx512{});
+        });
+    })
+    ONEDAL_IF_CPU_DISPATCH_AVX2(if (test_cpu_extension(cpu_ex, cpu_extension::avx2)) {
+        return task_executor_.execute([&]() {
+            return op(cpu_dispatch_avx2{});
+        });
+    })
+    ONEDAL_IF_CPU_DISPATCH_SSE42(if (test_cpu_extension(cpu_ex, cpu_extension::sse42)) {
+        return task_executor_.execute([&]() {
+            return op(cpu_dispatch_sse42{});
+        });
     })
-    ONEDAL_IF_CPU_DISPATCH_AVX2(
-        if (test_cpu_extension(cpu_ex, cpu_extension::avx2)) { return op(cpu_dispatch_avx2{}); })
-    ONEDAL_IF_CPU_DISPATCH_SSE42(
-        if (test_cpu_extension(cpu_ex, cpu_extension::sse42)) { return op(cpu_dispatch_sse42{}); })
-    return op(cpu_dispatch_default{});
+    return task_executor_.execute([&]() {
+        return op(cpu_dispatch_default{});
+    });
 }
 
 template <typename Op, typename OnUnknown>

cpp/oneapi/dal/backend/threading.cpp

@@ -0,0 +1,85 @@
+/*******************************************************************************
+* Copyright 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.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*******************************************************************************/
+#define TBB_PREVIEW_TASK_ARENA_CONSTRAINTS_EXTENSION 1
+
+#include "oneapi/dal/backend/threading.hpp"
+#include "src/threading/service_thread_pinner.h"
+#include "src/services/service_topo.h"
+
+namespace oneapi::dal::backend {
+
+struct task_daal : public daal::services::internal::thread_pinner_task_t {
+    task_daal(task& f) : task_oneapi(f) {}
+    void operator()() final {
+        task_oneapi();
+    }
+    task& task_oneapi;
+};
+
+class thread_pinner_impl {
+public:
+    thread_pinner_impl() {
+        daal_thread_pinner = std::shared_ptr<daal::services::internal::thread_pinner_t>(
+            daal::services::internal::getThreadPinner(true, read_topology, delete_topology));
+    }
+    thread_pinner_impl(tbb::task_arena& task_arena) {
+        daal_thread_pinner = std::shared_ptr<daal::services::internal::thread_pinner_t>(
+            daal::services::internal::getThreadPinnerFromTaskArena(true,
+                                                                   read_topology,
+                                                                   delete_topology,
+                                                                   task_arena));
+    }
+    void execute(task& task_) {
+        auto task_daal_ = task_daal(task_);
+        daal_thread_pinner->execute(task_daal_);
+    }
+
+private:
+    std::shared_ptr<daal::services::internal::thread_pinner_t> daal_thread_pinner;
+};
+
+tbb::task_arena* task_executor::create_task_arena(
+    const oneapi::dal::detail::threading_policy& policy) {
+    if (policy.max_threads_per_core) {
+        static tbb::task_arena task_arena{ tbb::task_arena::constraints{}.set_max_threads_per_core(
+            policy.max_threads_per_core) };
+        if (policy.thread_pinning) {
+            using daal::services::internal::thread_pinner_t;
+            thread_pinner_t* thread_pinner_ptr =
+                daal::services::internal::getThreadPinnerFromTaskArena(true,
+                                                                       read_topology,
+                                                                       delete_topology,
+                                                                       task_arena);
+            return thread_pinner_ptr->get_task_arena();
+        }
+        return &task_arena;
+    }
+    else {
+        static tbb::task_arena task_arena{};
+        if (policy.thread_pinning) {
+            using daal::services::internal::thread_pinner_t;
+            thread_pinner_t* thread_pinner_ptr =
+                daal::services::internal::getThreadPinnerFromTaskArena(true,
+                                                                       read_topology,
+                                                                       delete_topology,
+                                                                       task_arena);
+            return thread_pinner_ptr->get_task_arena();
+        }
+        return &task_arena;
+    }
+}
+
+} // namespace oneapi::dal::backend