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))
36typedef 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
52namespace tvm {
53namespace runtime {
54namespace threading {
55
56/*!
57 * \brief A platform-agnostic abstraction for managing a collection of
58 * thread pool threads.
59 */
60class 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 */
116void Yield();
117/*!
118 * \return the maximum number of effective workers for this system.
119 */
120int MaxConcurrency();
121/*!
122 * \brief Setting the maximum number of available cores.
123 */
124void 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 */
131void 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 */
140TVM_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 */
147int32_t NumThreads();
148
149} // namespace threading
150} // namespace runtime
151} // namespace tvm
152
153#endif // TVM_RUNTIME_THREADING_BACKEND_H_
154