1 // pthreadpool header from https://github.com/Maratyszcza/pthreadpool 2 // for NNPACK 3 #ifndef CAFFE2_UTILS_PTHREADPOOL_H_ 4 #define CAFFE2_UTILS_PTHREADPOOL_H_ 5 6 #include "ThreadPoolCommon.h" 7 8 #include <stddef.h> // for size_t 9 #include <stdint.h> // for uint32_t 10 11 #if defined(USE_PTHREADPOOL) 12 // This is a hack. 13 // Mainly introduced here because 14 // 1. NNPACK can be compiled to use internal legacy threadpool implementation because much of C2 depends on that. 15 // 2. Then if we want to use NNPACK in PyTorch, which uses new pthreadpool, then we will supply new pthreadpool pointer 16 // to NNPACK. This will not work if NNPACK is compiled with internal legacy threadpool. Thus this guard 17 // along with changes in pthreadpool_impl.cc allows us to override that behavior. 18 // It enables us to use NNPACK from pytorch using `caffe2::pthreadpool_()` 19 namespace caffe2 { 20 class WithCastToNewThreadPool { 21 public: 22 explicit WithCastToNewThreadPool(bool use_new_threadpool); 23 ~WithCastToNewThreadPool(); 24 private: 25 bool use_new_threadpool_; 26 }; 27 } 28 #endif 29 30 typedef struct pthreadpool* legacy_pthreadpool_t; 31 32 typedef void (*legacy_pthreadpool_function_1d_t)(void*, size_t); 33 typedef void (*legacy_pthreadpool_function_1d_tiled_t)(void*, size_t, size_t); 34 typedef void (*legacy_pthreadpool_function_2d_t)(void*, size_t, size_t); 35 typedef void (*legacy_pthreadpool_function_2d_tiled_t)(void*, size_t, size_t, size_t, size_t); 36 typedef void (*legacy_pthreadpool_function_3d_tiled_t)( 37 void*, 38 size_t, 39 size_t, 40 size_t, 41 size_t, 42 size_t, 43 size_t); 44 typedef void (*legacy_pthreadpool_function_4d_tiled_t)( 45 void*, 46 size_t, 47 size_t, 48 size_t, 49 size_t, 50 size_t, 51 size_t, 52 size_t, 53 size_t); 54 55 #ifdef __cplusplus 56 extern "C" { 57 #endif 58 59 /** 60 * Creates a thread pool with the specified number of threads. 61 * 62 * @param[in] threads_count The number of threads in the thread pool. 63 * A value of 0 has special interpretation: it creates a thread for each 64 * processor core available in the system. 65 * 66 * @returns A pointer to an opaque thread pool object. 67 * On error the function returns NULL and sets errno accordingly. 68 */ 69 70 // Returns internal threadpool impl. 71 legacy_pthreadpool_t legacy_pthreadpool_create(size_t threads_count); 72 73 /** 74 * Queries the number of threads in a thread pool. 75 * 76 * @param[in] threadpool The thread pool to query. 77 * 78 * @returns The number of threads in the thread pool. 79 */ 80 size_t legacy_pthreadpool_get_threads_count(legacy_pthreadpool_t threadpool); 81 82 /** 83 * Processes items in parallel using threads from a thread pool. 84 * 85 * When the call returns, all items have been processed and the thread pool is 86 * ready for a new task. 87 * 88 * @note If multiple threads call this function with the same thread pool, the 89 * calls are serialized. 90 * 91 * @param[in] threadpool The thread pool to use for parallelisation. 92 * @param[in] function The function to call for each item. 93 * @param[in] argument The first argument passed to the @a function. 94 * @param[in] items The number of items to process. The @a function 95 * will be called once for each item. 96 */ 97 void legacy_pthreadpool_compute_1d( 98 legacy_pthreadpool_t threadpool, 99 legacy_pthreadpool_function_1d_t function, 100 void* argument, 101 size_t range); 102 103 void legacy_pthreadpool_parallelize_1d( 104 legacy_pthreadpool_t threadpool, 105 legacy_pthreadpool_function_1d_t function, 106 void* argument, 107 size_t range, 108 uint32_t flags); 109 110 void legacy_pthreadpool_compute_1d_tiled( 111 legacy_pthreadpool_t threadpool, 112 legacy_pthreadpool_function_1d_tiled_t function, 113 void* argument, 114 size_t range, 115 size_t tile); 116 117 void legacy_pthreadpool_compute_2d( 118 legacy_pthreadpool_t threadpool, 119 legacy_pthreadpool_function_2d_t function, 120 void* argument, 121 size_t range_i, 122 size_t range_j); 123 124 void legacy_pthreadpool_compute_2d_tiled( 125 legacy_pthreadpool_t threadpool, 126 legacy_pthreadpool_function_2d_tiled_t function, 127 void* argument, 128 size_t range_i, 129 size_t range_j, 130 size_t tile_i, 131 size_t tile_j); 132 133 void legacy_pthreadpool_compute_3d_tiled( 134 legacy_pthreadpool_t threadpool, 135 legacy_pthreadpool_function_3d_tiled_t function, 136 void* argument, 137 size_t range_i, 138 size_t range_j, 139 size_t range_k, 140 size_t tile_i, 141 size_t tile_j, 142 size_t tile_k); 143 144 void legacy_pthreadpool_compute_4d_tiled( 145 legacy_pthreadpool_t threadpool, 146 legacy_pthreadpool_function_4d_tiled_t function, 147 void* argument, 148 size_t range_i, 149 size_t range_j, 150 size_t range_k, 151 size_t range_l, 152 size_t tile_i, 153 size_t tile_j, 154 size_t tile_k, 155 size_t tile_l); 156 157 /** 158 * Terminates threads in the thread pool and releases associated resources. 159 * 160 * @warning Accessing the thread pool after a call to this function constitutes 161 * undefined behaviour and may cause data corruption. 162 * 163 * @param[in,out] threadpool The thread pool to destroy. 164 */ 165 void legacy_pthreadpool_destroy(legacy_pthreadpool_t threadpool); 166 167 #ifdef USE_INTERNAL_PTHREADPOOL_IMPL 168 169 #define pthreadpool_t legacy_pthreadpool_t 170 #define pthreadpool_function_1d_t legacy_pthreadpool_function_1d_t 171 #define pthreadpool_function_1d_tiled_t legacy_pthreadpool_function_1d_tiled_t 172 #define pthreadpool_function_2d_t legacy_pthreadpool_function_2d_t 173 #define pthreadpool_function_2d_tiled_t legacy_pthreadpool_function_2d_tiled_t 174 #define pthreadpool_function_3d_tiled_t legacy_pthreadpool_function_3d_tiled_t 175 #define pthreadpool_function_4d_tiled_t legacy_pthreadpool_function_4d_tiled_t 176 #define pthreadpool_create legacy_pthreadpool_create 177 #define pthreadpool_destroy legacy_pthreadpool_destroy 178 #define pthreadpool_get_threads_count legacy_pthreadpool_get_threads_count 179 #define pthreadpool_compute_1d legacy_pthreadpool_compute_1d 180 #define pthreadpool_parallelize_1d legacy_pthreadpool_parallelize_1d 181 #define pthreadpool_compute_1d_tiled legacy_pthreadpool_compute_1d_tiled 182 #define pthreadpool_compute_2d legacy_pthreadpool_compute_2d 183 #define pthreadpool_compute_2d_tiled legacy_pthreadpool_compute_2d_tiled 184 #define pthreadpool_compute_3d_tiled legacy_pthreadpool_compute_3d_tiled 185 #define pthreadpool_compute_4d_tiled legacy_pthreadpool_compute_4d_tiled 186 187 #endif /* USE_INTERNAL_PTHREADPOOL_IMPL */ 188 189 #ifdef __cplusplus 190 } /* extern "C" */ 191 #endif 192 193 #endif // CAFFE2_UTILS_PTHREADPOOL_H_ 194