1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file tvm/runtime/threading_backend.h |
22 | * \brief Utilities for manipulating thread pool threads. |
23 | */ |
24 | #ifndef TVM_RUNTIME_THREADING_BACKEND_H_ |
25 | #define TVM_RUNTIME_THREADING_BACKEND_H_ |
26 | |
27 | #include <functional> |
28 | #include <memory> |
29 | #include <vector> |
30 | |
31 | #if defined(__linux__) || defined(__ANDROID__) |
32 | #if defined(__ANDROID__) |
33 | #ifndef CPU_SET |
34 | #define CPU_SETSIZE 1024 |
35 | #define __NCPUBITS (8 * sizeof(uint64_t)) |
36 | typedef struct { |
37 | uint64_t __bits[CPU_SETSIZE / __NCPUBITS]; |
38 | } cpu_set_t; |
39 | |
40 | #define CPU_SET(cpu, cpusetp) \ |
41 | ((cpusetp)->__bits[(cpu) / __NCPUBITS] |= (1UL << ((cpu) % __NCPUBITS))) |
42 | #define CPU_ZERO(cpusetp) memset((cpusetp), 0, sizeof(cpu_set_t)) |
43 | #define CPU_ISSET(cpu, cpusetp) \ |
44 | (1UL << ((cpu) % __NCPUBITS)) == \ |
45 | ((cpusetp)->__bits[(cpu) / __NCPUBITS] & (1UL << ((cpu) % __NCPUBITS))) |
46 | #define CPU_EQUAL(left, right) (memcmp(&left, &right, sizeof(cpu_set_t)) == 0) |
47 | |
48 | #endif |
49 | #endif |
50 | #endif |
51 | |
52 | namespace tvm { |
53 | namespace runtime { |
54 | namespace threading { |
55 | |
56 | /*! |
57 | * \brief A platform-agnostic abstraction for managing a collection of |
58 | * thread pool threads. |
59 | */ |
60 | class ThreadGroup { |
61 | public: |
62 | class Impl; |
63 | |
64 | /*! |
65 | * \brief Creates a collection of threads which run a provided function. |
66 | * |
67 | * \param num_workers The total number of worker threads in this group. |
68 | Includes main thread if `exclude_worker0 = true` |
69 | * \param worker_callback A callback which is run in its own thread. |
70 | Receives the worker_id as an argument. |
71 | * \param exclude_worker0 Whether to use the main thread as a worker. |
72 | * If `true`, worker0 will not be launched in a new thread and |
73 | * `worker_callback` will only be called for values >= 1. This |
74 | * allows use of the main thread as a worker. |
75 | */ |
76 | ThreadGroup(int num_workers, std::function<void(int)> worker_callback, |
77 | bool exclude_worker0 = false); |
78 | ~ThreadGroup(); |
79 | |
80 | /*! |
81 | * \brief Blocks until all non-main threads in the pool finish. |
82 | */ |
83 | void Join(); |
84 | |
85 | enum AffinityMode : int { |
86 | kBig = 1, |
87 | kLittle = -1, |
88 | /*Different threads will get different affinities.*/ |
89 | kSpecifyOneCorePerThread = -2, |
90 | /*All threads will get the same core group affinity.*/ |
91 | kSpecifyThreadShareAllCore = -3, |
92 | }; |
93 | /*! |
94 | * \brief configure the CPU id affinity |
95 | * |
96 | * \param mode The preferred CPU type (1 = big, -1 = little ...). |
97 | * \param nthreads The number of threads to use (0 = use all). |
98 | * \param exclude_worker0 Whether to use the main thread as a worker. |
99 | * If `true`, worker0 will not be launched in a new thread and |
100 | * `worker_callback` will only be called for values >= 1. This |
101 | * allows use of the main thread as a worker. |
102 | * \param cpus A list of CPU used to set 'cpu affinity'. |
103 | * |
104 | * \return The number of workers to use. |
105 | */ |
106 | int Configure(AffinityMode mode, int nthreads, bool exclude_worker0, |
107 | std::vector<unsigned int> cpus = {}); |
108 | |
109 | private: |
110 | Impl* impl_; |
111 | }; |
112 | |
113 | /*! |
114 | * \brief Platform-agnostic no-op. |
115 | */ |
116 | void Yield(); |
117 | /*! |
118 | * \return the maximum number of effective workers for this system. |
119 | */ |
120 | int MaxConcurrency(); |
121 | /*! |
122 | * \brief Setting the maximum number of available cores. |
123 | */ |
124 | void SetMaxConcurrency(int value); |
125 | /*! |
126 | * \brief Reset the threads in the pool. All current threads are destroyed and |
127 | * new ones are created. |
128 | * |
129 | * Note that this does nothing when openmp is used. |
130 | */ |
131 | void ResetThreadPool(); |
132 | |
133 | /*! |
134 | * \brief Configuring the CPU affinity mode for the working threads. |
135 | * \param mode The preferred CPU type (1 = big, -1 = little, -2 = kSpecifyOneCorePerThread, |
136 | * -3 = kSpecifyThreadShareAllCore). |
137 | * \param nthreads The number of threads to use (0 = use all). |
138 | * \param cpus A list of CPUs is used to set the 'cpu affinity' for the worker threads. |
139 | */ |
140 | TVM_DLL void Configure(tvm::runtime::threading::ThreadGroup::AffinityMode mode, int nthreads, |
141 | std::vector<unsigned int> cpus); |
142 | |
143 | /*! |
144 | * \brief Get the number of threads being used by the TVM runtime |
145 | * \returns The number of threads used. |
146 | */ |
147 | int32_t NumThreads(); |
148 | |
149 | } // namespace threading |
150 | } // namespace runtime |
151 | } // namespace tvm |
152 | |
153 | #endif // TVM_RUNTIME_THREADING_BACKEND_H_ |
154 | |