Pull request "Cache TBB local executors" by @NEvOlll from #2203

MERGED FROM #2203

ref:caae2e0d031d197a096970e90386b6cdcb19f7f8

catboost/python-package/catboost/_catboost.pyx

@@ -55,7 +55,7 @@ from util.generic.array_ref cimport TArrayRef, TConstArrayRef
 from util.generic.hash cimport THashMap
 from util.generic.hash_set cimport THashSet
 from util.generic.maybe cimport TMaybe
-from util.generic.ptr cimport THolder, TIntrusivePtr, MakeHolder
+from util.generic.ptr cimport TAtomicSharedPtr, THolder, TIntrusivePtr, MakeHolder
 from util.generic.string cimport TString, TStringBuf
 from util.generic.vector cimport TVector
 from util.system.types cimport ui8, ui16, ui32, ui64, i32, i64
@@ -410,13 +410,13 @@ cdef extern from "catboost/private/libs/options/load_options.h" namespace "NCatb
 
 cdef class Py_ObjectsOrderBuilderVisitor:
     cdef TDataProviderBuilderOptions options
-    cdef THolder[TTbbLocalExecutor] local_executor
+    cdef TAtomicSharedPtr[TTbbLocalExecutor] local_executor
     cdef THolder[IDataProviderBuilder] data_provider_builder
     cdef IRawObjectsOrderDataVisitor* builder_visitor
     cdef const TFeaturesLayout* features_layout
 
     def __cinit__(self, int thread_count):
-        self.local_executor = MakeHolder[TTbbLocalExecutor](thread_count)
+        self.local_executor = GetCachedLocalExecutor(thread_count)
         CreateDataProviderBuilderAndVisitor(
             self.options,
             <ILocalExecutor*>self.local_executor.Get(),
@@ -439,13 +439,13 @@ cdef class Py_ObjectsOrderBuilderVisitor:
 
 cdef class Py_FeaturesOrderBuilderVisitor:
     cdef TDataProviderBuilderOptions options
-    cdef THolder[TTbbLocalExecutor] local_executor
+    cdef TAtomicSharedPtr[TTbbLocalExecutor] local_executor
     cdef THolder[IDataProviderBuilder] data_provider_builder
     cdef IRawFeaturesOrderDataVisitor* builder_visitor
     cdef const TFeaturesLayout* features_layout
 
     def __cinit__(self, int thread_count):
-        self.local_executor = MakeHolder[TTbbLocalExecutor](thread_count)
+        self.local_executor = GetCachedLocalExecutor(thread_count)
         CreateDataProviderBuilderAndVisitor(
             self.options,
             <ILocalExecutor*>self.local_executor.Get(),
@@ -979,6 +979,7 @@ cdef extern from "catboost/python-package/catboost/helpers.h":
         int threadCount,
         ui64 cpuUsedRamLimit
     ) except +ProcessException
+    cdef TAtomicSharedPtr[TTbbLocalExecutor] GetCachedLocalExecutor(int threadsCount)
 
 
 cdef extern from "catboost/python-package/catboost/helpers.h":
@@ -4239,7 +4240,7 @@ cdef class _PoolBase:
         feature matrix : np.ndarray of shape (object_count, feature_count)
         """
         cdef int thread_count = UpdateThreadCount(-1)
-        cdef THolder[TTbbLocalExecutor] local_executor = MakeHolder[TTbbLocalExecutor](thread_count)
+        cdef TAtomicSharedPtr[TTbbLocalExecutor] local_executor = GetCachedLocalExecutor(thread_count)
         cdef TFeaturesLayout* features_layout =self.__pool.Get()[0].MetaInfo.FeaturesLayout.Get()
         cdef TRawObjectsDataProvider* raw_objects_data_provider = dynamic_cast_to_TRawObjectsDataProvider(
             self.__pool.Get()[0].ObjectsData.Get()
@@ -5621,7 +5622,7 @@ cdef class _StagedPredictIterator:
     cdef TVector[TVector[double]] __approx
     cdef TVector[TVector[double]] __pred
     cdef TFullModel* __model
-    cdef THolder[TTbbLocalExecutor] __executor
+    cdef TAtomicSharedPtr[TTbbLocalExecutor] __executor
     cdef TModelCalcerOnPool* __modelCalcerOnPool
     cdef EPredictionType predictionType
     cdef int ntree_start, ntree_end, eval_period, thread_count
@@ -5634,7 +5635,7 @@ cdef class _StagedPredictIterator:
         self.eval_period = eval_period
         self.thread_count = UpdateThreadCount(thread_count)
         self.verbose = verbose
-        self.__executor = MakeHolder[TTbbLocalExecutor](thread_count)
+        self.__executor = GetCachedLocalExecutor(thread_count)
 
     cdef _initialize_model_calcer(self, TFullModel* model, _PoolBase pool):
         self.__model = model

catboost/python-package/catboost/helpers.cpp

@@ -10,6 +10,10 @@
 #include <catboost/private/libs/options/split_params.h>
 #include <catboost/private/libs/target/data_providers.h>
 
+#include <util/system/guard.h>
+#include <util/system/info.h>
+#include <util/system/mutex.h>
+
 
 extern "C" PyObject* PyCatboostExceptionType;
 
@@ -295,3 +299,25 @@ void TrainEvalSplit(
         *evalDataProvider = getSubset(postShuffleTestIndices);
     }
 }
+
+TAtomicSharedPtr<NPar::TTbbLocalExecutor<false>> GetCachedLocalExecutor(int threadsCount) {
+    static TMutex lock;
+    static TAtomicSharedPtr<NPar::TTbbLocalExecutor<false>> cachedExecutor;
+
+    CB_ENSURE(threadsCount == -1 || 0 < threadsCount, "threadsCount should be positive or -1");    
+
+    if (threadsCount == -1) {
+        threadsCount = NSystemInfo::CachedNumberOfCpus();
+    }
+
+    with_lock (lock) {        
+        if (cachedExecutor && cachedExecutor->GetThreadCount() + 1 == threadsCount) {               
+            return cachedExecutor;
+        }                
+
+        cachedExecutor.Reset();
+        cachedExecutor = MakeAtomicShared<NPar::TTbbLocalExecutor<false>>(threadsCount);                
+
+        return cachedExecutor;
+    }    
+}

catboost/python-package/catboost/helpers.h

@@ -16,10 +16,12 @@
 #include <catboost/libs/train_lib/options_helper.h>
 
 #include <library/cpp/json/json_value.h>
+#include <library/cpp/threading/local_executor/tbb_local_executor.h>
 
 #include <util/generic/array_ref.h>
 #include <util/generic/fwd.h>
 #include <util/generic/noncopyable.h>
+#include <util/generic/ptr.h>
 #include <util/generic/xrange.h>
 
 #include <type_traits>
@@ -277,3 +279,4 @@ void TrainEvalSplit(
 );
 
 
+TAtomicSharedPtr<NPar::TTbbLocalExecutor<false>> GetCachedLocalExecutor(int threadsCount);