frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Re-apply PyTorch pthreadpool changes

Browse Source 
Summary:
This re-applies D21232894 (b9d3869df3) and D22162524, plus updates jni_deps in a few places
to avoid breaking host JNI tests.

Test Plan: `buck test @//fbandroid/mode/server //fbandroid/instrumentation_tests/com/facebook/caffe2:host-test`

Reviewed By: xcheng16

Differential Revision: D22199952

fbshipit-source-id: df13eef39c01738637ae8cf7f581d6ccc88d37d5
This commit is contained in:
David Reiss
2020-06-23 19:22:31 -07:00
committed by Facebook GitHub Bot
parent bdc00196d1
commit b7e044f0e5
44 changed files with 366 additions and 1756 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
caffe2/utils/threadpool/pthreadpool_impl.cc
View File

@ -6,9 +6,9 @@
// External API
//
void pthreadpool_compute_1d(
pthreadpool_t threadpool,
pthreadpool_function_1d_t function,
void legacy_pthreadpool_compute_1d(
legacy_pthreadpool_t threadpool,
legacy_pthreadpool_function_1d_t function,
void* argument,
size_t range) {
if (threadpool == nullptr) {
@ -27,30 +27,31 @@ void pthreadpool_compute_1d(
range);
}
size_t pthreadpool_get_threads_count(pthreadpool_t threadpool) {
// The current fix only useful when XNNPACK calls pthreadpool_get_threads_count with nullptr.
void legacy_pthreadpool_parallelize_1d(
const legacy_pthreadpool_t threadpool,
const legacy_pthreadpool_function_1d_t function,
void* const argument,
const size_t range,
uint32_t) {
legacy_pthreadpool_compute_1d(threadpool, function, argument, range);
}
size_t legacy_pthreadpool_get_threads_count(legacy_pthreadpool_t threadpool) {
// The current fix only useful when XNNPACK calls legacy_pthreadpool_get_threads_count with nullptr.
if (threadpool == nullptr) {
return 1;
}
return reinterpret_cast<caffe2::ThreadPool*>(threadpool)->getNumThreads();
// TODO: Future fix: If we keep maintaining two different threadpools.
// Old C2 and new one for XNNPACK, then the we have two different pthreadpool pointer
// types. One is caffe2::Thredpool*, the other is pthreadpool* (pthreadpool_new_if_impl.c)
// XNNPACK calls pthreadpool_get_threads_count during op setup using pthreadpool*, and
// uses _parallelize_ interface for for actual work.
// While NNPACK uses caffe2::Threadpool*.
// Thus if pthreadpool_get_threads_count is getting called from XNNPACK we cannot
// reinterpret_cast it to ThreadPool. It will seg fault or worse will have unedfined behavior.
}
pthreadpool_t pthreadpool_create(size_t threads_count) {
legacy_pthreadpool_t legacy_pthreadpool_create(size_t threads_count) {
std::mutex thread_pool_creation_mutex_;
std::lock_guard<std::mutex> guard(thread_pool_creation_mutex_);
return reinterpret_cast<pthreadpool_t>(new caffe2::ThreadPool(threads_count));
return reinterpret_cast<legacy_pthreadpool_t>(new caffe2::ThreadPool(threads_count));
}
void pthreadpool_destroy(pthreadpool_t pthreadpool) {
void legacy_pthreadpool_destroy(legacy_pthreadpool_t pthreadpool) {
if (pthreadpool) {
caffe2::ThreadPool* threadpool =
reinterpret_cast<caffe2::ThreadPool*>(pthreadpool);