File size: 2,293 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
#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
|