|
#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 |
|
} |
|
|
|
} |
|
|