1 %pythonbegin %{ 2 # 3 # Copyright (c) 2016-2021 Intel Corporation 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 18 __all__ = ["task_arena", 19 "task_group", 20 "global_control", 21 "default_num_threads", 22 "this_task_arena_max_concurrency", 23 "this_task_arena_current_thread_index", 24 "runtime_version", 25 "runtime_interface_version"] 26 %} 27 %begin %{ 28 /* Defines Python wrappers for Intel(R) oneAPI Threading Building Blocks (oneTBB) */ 29 %} 30 %module api 31 32 #if SWIG_VERSION < 0x030001 33 #error SWIG version 3.0.6 or newer is required for correct functioning 34 #endif 35 36 %{ 37 #define TBB_PREVIEW_WAITING_FOR_WORKERS 1 38 #include "tbb/task_arena.h" 39 #include "tbb/task_group.h" 40 #include "tbb/global_control.h" 41 #include "tbb/version.h" 42 43 #include <condition_variable> 44 #include <mutex> 45 #include <memory> 46 47 using namespace tbb; 48 49 class PyCaller : public swig::SwigPtr_PyObject { 50 public: 51 // icpc 2013 does not support simple using SwigPtr_PyObject::SwigPtr_PyObject; 52 PyCaller(const PyCaller& s) : SwigPtr_PyObject(s) {} 53 PyCaller(PyObject *p, bool initial = true) : SwigPtr_PyObject(p, initial) {} 54 55 void operator()() const { 56 SWIG_PYTHON_THREAD_BEGIN_BLOCK; 57 PyObject* r = PyObject_CallFunctionObjArgs((PyObject*)*this, nullptr); 58 if(r) Py_DECREF(r); 59 SWIG_PYTHON_THREAD_END_BLOCK; 60 } 61 }; 62 63 struct ArenaPyCaller { 64 task_arena *my_arena; 65 PyObject *my_callable; 66 ArenaPyCaller(task_arena *a, PyObject *c) : my_arena(a), my_callable(c) { 67 SWIG_PYTHON_THREAD_BEGIN_BLOCK; 68 Py_XINCREF(c); 69 SWIG_PYTHON_THREAD_END_BLOCK; 70 } 71 void operator()() const { 72 my_arena->execute(PyCaller(my_callable, false)); 73 } 74 }; 75 76 struct barrier_data { 77 std::condition_variable event; 78 std::mutex m; 79 int worker_threads, full_threads; 80 }; 81 82 void _concurrency_barrier(int threads = tbb::task_arena::automatic) { 83 if(threads == tbb::task_arena::automatic) 84 threads = tbb::this_task_arena::max_concurrency(); 85 if(threads < 2) 86 return; 87 std::unique_ptr<global_control> g( 88 (global_control::active_value(global_control::max_allowed_parallelism) < unsigned(threads))? 89 new global_control(global_control::max_allowed_parallelism, threads) : nullptr); 90 91 tbb::task_group tg; 92 barrier_data b; 93 b.worker_threads = 0; 94 b.full_threads = threads-1; 95 for(int i = 0; i < b.full_threads; i++) 96 tg.run([&b]{ 97 std::unique_lock<std::mutex> lock(b.m); 98 if(++b.worker_threads >= b.full_threads) 99 b.event.notify_all(); 100 else while(b.worker_threads < b.full_threads) 101 b.event.wait(lock); 102 }); 103 std::unique_lock<std::mutex> lock(b.m); 104 b.event.wait(lock); 105 tg.wait(); 106 }; 107 108 %} 109 110 void _concurrency_barrier(int threads = tbb::task_arena::automatic); 111 112 namespace tbb { 113 114 class task_arena { 115 public: 116 static const int automatic = -1; 117 task_arena(int max_concurrency = automatic, unsigned reserved_for_masters = 1); 118 task_arena(const task_arena &s); 119 ~task_arena(); 120 void initialize(); 121 void initialize(int max_concurrency, unsigned reserved_for_masters = 1); 122 void terminate(); 123 bool is_active(); 124 %extend { 125 void enqueue( PyObject *c ) { $self->enqueue(PyCaller(c)); } 126 void execute( PyObject *c ) { $self->execute(PyCaller(c)); } 127 }; 128 }; 129 130 class task_group { 131 public: 132 task_group(); 133 ~task_group(); 134 void wait(); 135 void cancel(); 136 %extend { 137 void run( PyObject *c ) { $self->run(PyCaller(c)); } 138 void run( PyObject *c, task_arena *a ) { $self->run(ArenaPyCaller(a, c)); } 139 }; 140 }; 141 142 class global_control { 143 public: 144 enum parameter { 145 max_allowed_parallelism, 146 thread_stack_size, 147 parameter_max // insert new parameters above this point 148 }; 149 global_control(parameter param, size_t value); 150 ~global_control(); 151 static size_t active_value(parameter param); 152 }; 153 154 } // tbb 155 156 %inline { 157 inline const char* runtime_version() { return TBB_runtime_version();} 158 inline int runtime_interface_version() { return TBB_runtime_interface_version();} 159 inline int this_task_arena_max_concurrency() { return this_task_arena::max_concurrency();} 160 inline int this_task_arena_current_thread_index() { return this_task_arena::current_thread_index();} 161 }; 162 163 // Additional definitions for Python part of the module 164 %pythoncode %{ 165 default_num_threads = this_task_arena_max_concurrency 166 %} 167