forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Revert "[Submodule] Remove deprecated USE_TBB option and TBB submodule (
pytorch#127051)" This reverts commit 699db79. Reverted pytorch#127051 on behalf of https://github.com/PaliC due to This PR needs to be synced using the import button as there is a bug in our diff train ([comment](pytorch#127051 (comment)))
- Loading branch information
1 parent
1abcac9
commit 67739d8
Showing
34 changed files
with
863 additions
and
19 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,115 @@ | ||
| #include <ATen/Config.h> | ||
| #if AT_PARALLEL_NATIVE_TBB | ||
| #include <ATen/Parallel.h> | ||
| #include <ATen/ParallelFuture.h> | ||
| #include <ATen/PTThreadPool.h> | ||
|
|
||
| #include <atomic> | ||
| #include <mutex> | ||
|
|
||
| #include <tbb/tbb.h> | ||
| #define TBB_PREVIEW_GLOBAL_CONTROL 1 | ||
| #include <tbb/global_control.h> | ||
|
|
||
| #ifdef _OPENMP | ||
| #include <omp.h> | ||
| #endif | ||
|
|
||
| #if AT_MKL_ENABLED() | ||
| #include <mkl.h> | ||
| #endif | ||
|
|
||
| namespace at { | ||
|
|
||
| namespace { | ||
| static thread_local tbb::task_group tg_; | ||
| thread_local int this_thread_id{0}; | ||
|
|
||
| std::mutex global_thread_mutex_; | ||
| std::shared_ptr<tbb::global_control> global_thread_limit_ = nullptr; | ||
| std::atomic<int> num_intraop_threads_{-1}; | ||
|
|
||
| void _internal_set_num_threads(int nthreads) { | ||
| TORCH_INTERNAL_ASSERT(nthreads > 0); | ||
| { | ||
| std::unique_lock<std::mutex> lk(global_thread_mutex_); | ||
| // This is an antipattern and we shouldn't be constraining the number of | ||
| // threads in library code. | ||
| // TODO: Think of a smarter way to leverage tbb::thread_arena to limit the | ||
| // number of slots instead of the number of threads. | ||
| global_thread_limit_ = std::make_shared<tbb::global_control>( | ||
| tbb::global_control::max_allowed_parallelism, nthreads); | ||
| num_intraop_threads_.store(nthreads); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| void init_num_threads() { | ||
| #ifdef _OPENMP | ||
| omp_set_num_threads(1); | ||
| #endif | ||
|
|
||
| #if AT_MKL_ENABLED() | ||
| mkl_set_num_threads(1); | ||
| #endif | ||
|
|
||
| int nthreads = num_intraop_threads_.load(); | ||
| if (nthreads < 0) { | ||
| nthreads = intraop_default_num_threads(); | ||
| } | ||
| _internal_set_num_threads(nthreads); | ||
| } | ||
|
|
||
| void set_num_threads(int nthreads) { | ||
| TORCH_CHECK(nthreads > 0); | ||
|
|
||
| _internal_set_num_threads(nthreads); | ||
| } | ||
|
|
||
| int get_num_threads() { | ||
| at::internal::lazy_init_num_threads(); | ||
| return tbb::global_control::active_value( | ||
| tbb::global_control::max_allowed_parallelism); | ||
| } | ||
|
|
||
| int get_thread_num() { | ||
| return this_thread_id; | ||
| } | ||
|
|
||
| namespace internal { | ||
| void set_thread_num(int id) { | ||
| this_thread_id = id; | ||
| } | ||
| } | ||
|
|
||
| bool in_parallel_region() { | ||
| return tbb::this_task_arena::current_thread_index() >= 0; | ||
| } | ||
|
|
||
| void intraop_launch(std::function<void()> func) { | ||
| if (get_num_threads() > 1) { | ||
| tg_.run(func); | ||
| } else { | ||
| func(); | ||
| } | ||
| } | ||
|
|
||
| c10::intrusive_ptr<c10::ivalue::Future> intraop_launch_future( | ||
| std::function<void()> func) { | ||
| auto future = c10::make_intrusive<c10::ivalue::Future>(NoneType::get()); | ||
| if (get_num_threads() > 1) { | ||
| tg_.run( | ||
| [func, future]() { | ||
| func(); | ||
| future->markCompleted(); | ||
| } | ||
| ); | ||
| } else { | ||
| func(); | ||
| future->markCompleted(); | ||
| } | ||
| return future; | ||
| } | ||
|
|
||
| } // namespace at | ||
| #endif |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| #pragma once | ||
|
|
||
| #include <atomic> | ||
| #include <cstddef> | ||
| #include <exception> | ||
|
|
||
| #include <c10/util/Exception.h> | ||
|
|
||
| #ifdef _WIN32 | ||
| #ifndef WIN32_LEAN_AND_MEAN | ||
| #define WIN32_LEAN_AND_MEAN | ||
| #endif | ||
| #endif | ||
| #include <tbb/tbb.h> | ||
|
|
||
| #define INTRA_OP_PARALLEL | ||
|
|
||
| namespace at::internal { | ||
|
|
||
| template <typename F> | ||
| inline void invoke_parallel( | ||
| const int64_t begin, | ||
| const int64_t end, | ||
| const int64_t grain_size, | ||
| const F& f) { | ||
| // Choose number of tasks based on grain size and number of threads. | ||
| int64_t chunk_size = divup((end - begin), get_num_threads()); | ||
| // Make sure each task is at least grain_size size. | ||
| chunk_size = std::max(grain_size, chunk_size); | ||
|
|
||
| std::atomic_flag err_flag = ATOMIC_FLAG_INIT; | ||
| std::exception_ptr eptr; | ||
| tbb::parallel_for( | ||
| tbb::blocked_range<int64_t>(begin, end, chunk_size), | ||
| [&eptr, &err_flag, f](const tbb::blocked_range<int64_t>& r) { | ||
| try { | ||
| internal::ThreadIdGuard tid_guard( | ||
| tbb::this_task_arena::current_thread_index()); | ||
| f(r.begin(), r.end()); | ||
| } catch (...) { | ||
| if (!err_flag.test_and_set()) { | ||
| eptr = std::current_exception(); | ||
| } | ||
| } | ||
| }, | ||
| tbb::static_partitioner{}); | ||
| if (eptr) { | ||
| std::rethrow_exception(eptr); | ||
| } | ||
| } | ||
|
|
||
| } // namespace at::internal |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.