Changing threading control interfaces (WIP)

cpp/daal/BUILD

@@ -21,7 +21,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/include/services/env_detect.h

@@ -169,6 +169,12 @@ class DAAL_EXPORT Environment : public Base
 
     services::internal::sycl::ExecutionContextIface & getDefaultExecutionContext() { return *_executionContext; }
 
+    /**
+     *  Enables external threading control
+     *  \param[in] assumeExternalControl   Flag to assume external threading control
+     */
+    void modifyExternalThreadingControl(bool assumeExternalControl = true);
+
 private:
     Environment();
     Environment(const Environment & e);

cpp/daal/src/services/env_detect.cpp

@@ -33,6 +33,7 @@
 
 static daal::services::Environment::LibraryThreadingType daal_thr_set = (daal::services::Environment::LibraryThreadingType)-1;
 static bool isInit                                                    = false;
+static bool externalThreadingControl                                  = false;
 
 namespace daal
 {
@@ -126,21 +127,22 @@ DAAL_EXPORT daal::services::Environment::Environment(const Environment & e) : da
 
 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())
+    if (!isInit && !externalThreadingControl)
     {
-        /* 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()))
+        /* if HT enabled - set _numThreads to physical cores num */
+        if (daal::internal::ServiceInst::serv_get_ht())
         {
-            daal::services::Environment::setNumberOfThreads(ncores);
+            /* 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;
     }
-    isInit = true;
 }
 
 DAAL_EXPORT daal::services::Environment::~Environment()
@@ -160,7 +162,6 @@ void daal::services::Environment::_cpu_detect(int enable)
 
 DAAL_EXPORT void daal::services::Environment::setNumberOfThreads(const size_t numThreads)
 {
-    isInit = true;
     daal::setNumberOfThreads(numThreads, &_globalControl);
 }
 
@@ -188,3 +189,13 @@ DAAL_EXPORT void daal::services::Environment::enableThreadPinning(const bool ena
 #endif
     return;
 }
+
+DAAL_EXPORT void daal::services::Environment::modifyExternalThreadingControl(bool assumeExternalControl)
+{
+    externalThreadingControl = assumeExternalControl;
+    isInit                   = false;
+    if (externalThreadingControl)
+    {
+        daal::services::Environment::setNumberOfThreads(_daal_threader_get_max_threads());
+    }
+}

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.cpp

@@ -31,7 +31,11 @@ struct global_context_cpu_init {
 };
 
 void context_cpu::global_init() {
+    using daal::services::Environment;
+
     [[maybe_unused]] static volatile global_context_cpu_init init;
+
+    Environment::getInstance()->modifyExternalThreadingControl();
 }
 
 inline constexpr detail::cpu_extension from_daal_cpu_type(daal::CpuType cpu) {