-
Notifications
You must be signed in to change notification settings - Fork 3.7k
[Runtime][ThreadPool]Refactor affinity function and support CPU affinity list setting. #9802
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
a01df59
1b457e9
0c8d189
d8ec871
7922d6e
b3cdd2a
937aa77
ada242c
13fe147
770fc92
7ea0c2e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,6 +24,7 @@ | |
| #include <dmlc/thread_local.h> | ||
| #include <tvm/runtime/c_backend_api.h> | ||
| #include <tvm/runtime/c_runtime_api.h> | ||
| #include <tvm/runtime/container/array.h> | ||
| #include <tvm/runtime/logging.h> | ||
| #include <tvm/runtime/packed_func.h> | ||
| #include <tvm/runtime/registry.h> | ||
|
|
@@ -42,12 +43,13 @@ | |
| #include <thread> | ||
| #include <vector> | ||
|
|
||
| #include "../support/utils.h" | ||
| const constexpr int kL1CacheBytes = 64; | ||
|
|
||
| namespace tvm { | ||
| namespace runtime { | ||
| namespace { | ||
|
|
||
| using support::IsNumber; | ||
| constexpr uint32_t kDefaultSpinCount = 300000; | ||
|
|
||
| uint32_t GetSpinCount() { | ||
|
|
@@ -317,10 +319,11 @@ class ThreadPool { | |
|
|
||
| static ThreadPool* ThreadLocal() { return dmlc::ThreadLocalStore<ThreadPool>::Get(); } | ||
|
|
||
| void UpdateWorkerConfiguration(threading::ThreadGroup::AffinityMode mode, int nthreads) { | ||
| void UpdateWorkerConfiguration(threading::ThreadGroup::AffinityMode mode, int nthreads, | ||
| const std::vector<unsigned int>& cpus) { | ||
| // this will also reset the affinity of the ThreadGroup | ||
| // may use less than the MaxConcurrency number of workers | ||
| num_workers_used_ = threads_->Configure(mode, nthreads, exclude_worker0_); | ||
| num_workers_used_ = threads_->Configure(mode, nthreads, exclude_worker0_, cpus); | ||
| // if MaxConcurrency restricted the number of workers (e.g., due to | ||
| // hyperthreading), respect the restriction | ||
| num_workers_used_ = std::min(num_workers_, num_workers_used_); | ||
|
|
@@ -369,17 +372,42 @@ class ThreadPool { | |
| std::unique_ptr<tvm::runtime::threading::ThreadGroup> threads_; | ||
| }; | ||
|
|
||
| /*! | ||
| * \brief args[0] is the AffinityMode, args[1] is the number of threads. | ||
| * args2 is a list of CPUs which is used to set the CPU affinity. | ||
| */ | ||
| TVM_REGISTER_GLOBAL("runtime.config_threadpool").set_body([](TVMArgs args, TVMRetValue* rv) { | ||
| threading::ThreadGroup::AffinityMode mode = | ||
| static_cast<threading::ThreadGroup::AffinityMode>(static_cast<int>(args[0])); | ||
| int nthreads = args[1]; | ||
| ThreadPool::ThreadLocal()->UpdateWorkerConfiguration(mode, nthreads); | ||
| std::vector<unsigned int> cpus; | ||
| if (args.num_args >= 3) { | ||
| Array<String> cpu_array = args[2]; | ||
| for (auto cpu : cpu_array) { | ||
| ICHECK(IsNumber(cpu)) << "The CPU core information '" << cpu << "' is not a number."; | ||
| cpus.push_back(std::stoi(cpu)); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Verify that the string represents a valid integer?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What you want to do is to restrict tvm thread run on specific cpu core ids? If so, how to handle Python's API interface? For example: from tvm._ffi import get_global_func
config_threadpool = get_global_func('runtime.config_threadpool')
core_ids = (0, 1, 2)
config_threadpool(0, 1, *core_ids)?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. thanks @FrozenGene for the follow up, if the core ids is greater than threads number, all the threads will be set the affinity with all of the cpu in cpu list, at the said case, thread 0 will affinity with cpu (0, 1, 2) , the related logic in threading_backend.cc::120 - 129 line.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What I asked is how to handle Python's API pack syntax. If I write the code as previous, your current code can not handle, because the unpacked argument will not the type of list like C++. @huajsj
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sorry for the misunderstanding. following are the answer
Yes, restrict the worker thread running on specific cpu core or cpu core groups
the supported use case like following
In existing logic, the second parameter is not used to determine how many worker threads to launch, it is used as
|
||
| std::cout << "cpu is " << cpu << std::endl; | ||
| } | ||
| } | ||
| threading::Configure(mode, nthreads, cpus); | ||
| }); | ||
|
|
||
| namespace threading { | ||
| void ResetThreadPool() { tvm::runtime::ThreadPool::ThreadLocal()->Reset(); } | ||
| /*! | ||
| * \brief configure the CPU id affinity | ||
| * \param mode The preferred CPU type (1 = big, -1 = little, -2 = specify , | ||
| * -3 = kSpecifyOneCorePerThread, -3 = kSpecifyThreadShareAllCore). | ||
| * \param nthreads The number of threads to use (0 = use all). | ||
| * \param cpus cpus A list of CPUs is used to set the 'cpu affinity' for the worker threads. | ||
| * | ||
| */ | ||
| void Configure(tvm::runtime::threading::ThreadGroup::AffinityMode mode, int nthreads, | ||
| std::vector<unsigned int> cpus) { | ||
| tvm::runtime::threading::SetMaxConcurrency(cpus.size()); | ||
| tvm::runtime::ThreadPool::ThreadLocal()->UpdateWorkerConfiguration(mode, nthreads, cpus); | ||
| } | ||
| } // namespace threading | ||
|
|
||
| } // namespace runtime | ||
| } // namespace tvm | ||
|
|
||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.