#pragma once
|
#include <ATen/ATen.h>
|
#include <ATen/core/ivalue.h>
|
|
namespace at {
|
namespace internal {
|
// This parameter is heuristically chosen to determine the minimum number of
|
// work that warrants parallelism. For example, when summing an array, it is
|
// deemed inefficient to parallelise over arrays shorter than 32768. Further,
|
// no parallel algorithm (such as parallel_reduce) should split work into
|
// smaller than GRAIN_SIZE chunks.
|
constexpr int64_t GRAIN_SIZE = 32768;
|
} // namespace internal
|
|
inline int64_t divup(int64_t x, int64_t y) {
|
return (x + y - 1) / y;
|
}
|
|
// Called during new thread initialization
|
CAFFE2_API void init_num_threads();
|
|
// Sets the number of threads to be used in parallel region
|
CAFFE2_API void set_num_threads(int);
|
|
// Returns the number of threads used in parallel region
|
CAFFE2_API int get_num_threads();
|
|
// Returns the current thread number (starting from 0)
|
// in the current parallel region, or 0 in the sequential region
|
CAFFE2_API int get_thread_num();
|
|
// Checks whether the code runs in parallel region
|
CAFFE2_API bool in_parallel_region();
|
|
/*
|
parallel_for
|
|
begin: index at which to start applying user function
|
|
end: index at which to stop applying user function
|
|
grain_size: number of elements per chunk. impacts the degree of parallelization
|
|
f: user function applied in parallel to the chunks, signature:
|
void f(int64_t begin, int64_t end)
|
*/
|
template <class F>
|
inline void parallel_for(
|
const int64_t begin,
|
const int64_t end,
|
const int64_t grain_size,
|
const F& f);
|
|
/*
|
parallel_reduce
|
|
begin: index at which to start applying reduction
|
|
end: index at which to stop applying reduction
|
|
grain_size: number of elements per chunk. impacts number of elements in
|
intermediate results tensor and degree of parallelization.
|
|
ident: identity for binary combination function sf. sf(ident, x) needs to return
|
x.
|
|
f: function for reduction over a chunk. f needs to be of signature scalar_t
|
f(int64_t partial_begin, int64_t partial_end, scalar_t identifiy)
|
|
sf: function to combine two partial results. sf needs to be of signature
|
scalar_t sf(scalar_t x, scalar_t y)
|
|
For example, you might have a tensor of 10000 entires and want to sum together
|
all the elements. Parallel_reduce with a grain_size of 2500 will then allocate
|
an intermediate result tensor with 4 elements. Then it will execute the function
|
"f" you provide and pass the beginning and end index of these chunks, so
|
0-2499, 2500-4999, etc. and the combination identity. It will then write out
|
the result from each of these chunks into the intermediate result tensor. After
|
that it'll reduce the partial results from each chunk into a single number using
|
the combination function sf and the identity ident. For a total summation this
|
would be "+" and 0 respectively. This is similar to tbb's approach [1], where
|
you need to provide a function to accumulate a subrange, a function to combine
|
two partial results and an identity.
|
|
[1] https://software.intel.com/en-us/node/506154
|
*/
|
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);
|
|
// Returns a detailed string describing parallelization settings
|
CAFFE2_API std::string get_parallel_info();
|
|
// Sets number of threads used for inter-op parallelism
|
CAFFE2_API void set_num_interop_threads(int);
|
|
// Returns the number of threads used for inter-op parallelism
|
CAFFE2_API int get_num_interop_threads();
|
|
// Launches inter-op parallel task
|
CAFFE2_API void launch(std::function<void()> func);
|
|
// Launches intra-op parallel task
|
CAFFE2_API void intraop_launch(std::function<void()> func);
|
|
// Launches intra-op parallel task, returns a future
|
CAFFE2_API std::shared_ptr<c10::ivalue::Future> intraop_launch_future(
|
std::function<void()> func);
|
|
// Returns number of intra-op threads used by default
|
CAFFE2_API int intraop_default_num_threads();
|
|
} // namespace at
|
|
#if AT_PARALLEL_OPENMP
|
#include <ATen/ParallelOpenMP.h>
|
#elif AT_PARALLEL_NATIVE
|
#include <ATen/ParallelNative.h>
|
#elif AT_PARALLEL_NATIVE_TBB
|
#include <ATen/ParallelNativeTBB.h>
|
#endif
|
