94 lines
2.2 KiB
C
94 lines
2.2 KiB
C
|
#pragma once
|
||
|
|
|
||
|
|
#include <c10/util/Exception.h>
|
||
|
|
#include <c10/util/ParallelGuard.h>
|
||
|
|
#include <c10/util/SmallVector.h>
|
||
|
|
|
||
|
|
namespace at {
|
||
|
|
|
||
|
|
template <class F>
|
||
|
|
inline void parallel_for(
|
||
|
|
const int64_t begin,
|
||
|
|
const int64_t end,
|
||
|
|
const int64_t grain_size,
|
||
|
|
const F& f) {
|
||
|
|
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(grain_size >= 0);
|
||
|
|
if (begin >= end) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
#ifdef INTRA_OP_PARALLEL
|
||
|
|
at::internal::lazy_init_num_threads();
|
||
|
|
const auto numiter = end - begin;
|
||
|
|
const bool use_parallel =
|
||
|
|
(numiter > grain_size && numiter > 1 && !at::in_parallel_region() &&
|
||
|
|
at::get_num_threads() > 1);
|
||
|
|
if (!use_parallel) {
|
||
|
|
internal::ThreadIdGuard tid_guard(0);
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
f(begin, end);
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
internal::invoke_parallel(
|
||
|
|
begin, end, grain_size, [&](int64_t begin, int64_t end) {
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
f(begin, end);
|
||
|
|
});
|
||
|
|
#else
|
||
|
|
internal::ThreadIdGuard tid_guard(0);
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
f(begin, end);
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class scalar_t, class F, class SF>
|
||
|
|
inline scalar_t parallel_reduce(
|
||
|
|
const int64_t begin,
|
||
|
|
const int64_t end,
|
||
|
|
const int64_t grain_size,
|
||
|
|
const scalar_t ident,
|
||
|
|
const F& f,
|
||
|
|
const SF& sf) {
|
||
|
|
TORCH_CHECK(grain_size >= 0);
|
||
|
|
if (begin >= end) {
|
||
|
|
return ident;
|
||
|
|
}
|
||
|
|
|
||
|
|
#ifdef INTRA_OP_PARALLEL
|
||
|
|
at::internal::lazy_init_num_threads();
|
||
|
|
const auto max_threads = at::get_num_threads();
|
||
|
|
const bool use_parallel =
|
||
|
|
((end - begin) > grain_size && !at::in_parallel_region() &&
|
||
|
|
max_threads > 1);
|
||
|
|
if (!use_parallel) {
|
||
|
|
internal::ThreadIdGuard tid_guard(0);
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
return f(begin, end, ident);
|
||
|
|
}
|
||
|
|
|
||
|
|
c10::SmallVector<scalar_t, 64> results(max_threads, ident);
|
||
|
|
internal::invoke_parallel(
|
||
|
|
begin,
|
||
|
|
end,
|
||
|
|
grain_size,
|
||
|
|
[&](const int64_t my_begin, const int64_t my_end) {
|
||
|
|
const auto tid = at::get_thread_num();
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
results[tid] = f(my_begin, my_end, ident);
|
||
|
|
});
|
||
|
|
|
||
|
|
scalar_t result = ident;
|
||
|
|
for (auto partial_result : results) {
|
||
|
|
result = sf(result, partial_result);
|
||
|
|
}
|
||
|
|
return result;
|
||
|
|
#else
|
||
|
|
internal::ThreadIdGuard tid_guard(0);
|
||
|
|
c10::ParallelGuard guard(true);
|
||
|
|
return f(begin, end, ident);
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace at
|