saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve work distribution for Expand operator, and sharded LoopCounter configuration (#6454)

Browse source 
Description: This PR makes two changes identified while looking at a PGAN model.

First, it uses ThreadPool::TryParallelFor for the main parallel loops in the Expand operator. This lets the thread pool decide on the granularity at which to distribute work (unlike TrySimpleParallelFor). Profiling showed high costs when running "simple" loops with 4M iterations each of which copied only 4 bytes.

Second, it updates the sharded loop counter in the thread pool so that the number of shards is capped by the number of threads. This helps make the performance of any other high-contention "simple" loops more robust at low thread counts by letting each thread work on its own "home" shard for longer.

Motivation and Context

Profiling showed a PGAN model taking 2x+ longer with the non-OpenMP build. The root cause was that the OpenMP build uses simple static scheduling of loop iterations, while the non-OpenMP build uses dynamic scheduling. The combination of large numbers of tiny iterations is less significant with static scheduling --- although still desirable to avoid, given that each iteration incurs a std::function invocation.
This commit is contained in:
Tim Harris 2021-01-29 11:19:18 +00:00 committed by GitHub
parent 7abb5b667f
commit 066520f6c1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 59 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
onnxruntime/core/common/threadpool.cc
View file

@ -28,8 +28,8 @@ namespace concurrency {
// A sharded loop counter distributes loop iterations between a set of worker threads. The iteration space of
// the loop is divided (perhaps unevenly) between the shards. Each thread has a home shard (perhaps not uniquely
// to it), and it claims iterations via atomic operations on its home shard. It then proceeds through the other
// shards until all of the shards' iterations are complete. This approach serves to purposes. First, compared
// with atomic operations on a single counter, it reduces contention on a single counter in the case of loops with
// shards until all of the shards' iterations are complete. This approach serves two purposes. First, compared
// with atomic operations on a single counter, it reduces contention on the counter in the case of loops with
// large numbers of short-running iteration. Second, by having a thread work on its home shard initially, it
// promotes affinity between the work that a thread performs in one loop and the work that it performs in the next.
@ -51,8 +51,11 @@ static_assert(sizeof(LoopCounterShard) == CACHE_LINE_BYTES, "Expected loop count
class alignas(CACHE_LINE_BYTES) LoopCounter {
public:
LoopCounter(uint64_t num_iterations,
uint64_t d_of_p,
uint64_t block_size = 1) : _block_size(block_size),
_num_shards(GetNumShards(num_iterations, block_size)) {
_num_shards(GetNumShards(num_iterations,
d_of_p,
block_size)) {
// Divide the iteration space between the shards. If the iteration
// space does not divide evenly into shards of multiples of
// block_size then the final shard is left uneven.
@ -111,9 +114,18 @@ public:
private:
// Derive the number of shards to use for a given loop. We require
// at least one block of work per shard, and subject to that
// constraint we use [1,MAX_SHARDS) shards.
// at least one block of work per shard, and subject to the
// constraints:
//
// - We use no more than MAX_SHARDS (limiting the amount of space needed
// for the LoopCounter, and work needed to confirm that all shards have been
// completed at the end of a loop).
//
// - The number of shards is <= the number of threads (d_of_p).
// Hence, at low thread counts, each of N threads will get its own
// shard representing 1/N of the work.
static unsigned GetNumShards(uint64_t num_iterations,
uint64_t d_of_p,
uint64_t block_size) {
unsigned num_shards;
auto num_blocks = num_iterations / block_size;
@ -124,6 +136,9 @@ private:
} else {
num_shards = MAX_SHARDS;
}
if (num_shards > d_of_p) {
num_shards = static_cast<unsigned>(d_of_p);
}
return num_shards;
}
@ -181,7 +196,7 @@ void ThreadPool::ParallelForFixedBlockSizeScheduling(const std::ptrdiff_t total,
int num_work_items = static_cast<int>(std::min(static_cast<std::ptrdiff_t>(d_of_p), num_blocks));
assert(num_work_items > 0);
LoopCounter lc(total, block_size);
LoopCounter lc(total, d_of_p, block_size);
std::function<void(unsigned)> run_work = [&](unsigned idx) {
unsigned my_home_shard = lc.GetHomeShard(idx);
unsigned my_shard = my_home_shard;

74
onnxruntime/core/providers/cpu/tensor/expand.cc
View file

@ -115,68 +115,70 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
auto copy_byte = copy_len * sizeof(T);
auto distribute_fn =
[&](ptrdiff_t i) {
[&](ptrdiff_t i_start, ptrdiff_t i_end) {
for (auto i = i_start; i < i_end; i++) {
auto input_offset = i * copy_len;
int64_t output_offset = 0;
for (auto j = dim_group_start + 1, remains = input_offset; j < max_dims_size; ++j) {
auto current_count = remains / input_dim_group[j];
output_offset += current_count * output_dim_group[j];
remains = remains % input_dim_group[j];
} //for
} //for j
memcpy(output_data + output_offset, input_data + input_offset, copy_byte);
output_offsets[i] = output_offset;
}; //distribute_fn
} //for i
}; //distribute_fn
auto per_thread_tasks =
distribute_count / concurrency::ThreadPool::DegreeOfParallelism(context->GetOperatorThreadPool());
if (per_thread_tasks > 4) {
concurrency::ThreadPool::TrySimpleParallelFor(
concurrency::ThreadPool::TryParallelFor(
context->GetOperatorThreadPool(),
distribute_count,
static_cast<double>(copy_byte),
std::move(distribute_fn));
} else {
for (int64_t i = 0; i < distribute_count; ++i) {
distribute_fn(i);
}
distribute_fn(0, distribute_count);
} //else
for (auto i = max_dims_size - 1; i >= dim_group_start; --i) {
auto copy_fn =
[&](ptrdiff_t j) {
auto output_offset = output_offsets[j];
if (output_offset % output_dim_group[i] == 0) {
auto copy_len = output_dim_group[i] / expand_dim_size[i];
auto copy_byte = copy_len * sizeof(T);
auto output_from = output_data + output_offset;
auto output_at = output_from + copy_len;
auto output_end = output_from + output_dim_group[i];
while (output_at + copy_len <= output_end) {
memcpy(output_at, output_from, copy_byte);
output_at += copy_len;
copy_len <<= 1;
copy_byte <<= 1;
} //while
while (output_at < output_end) {
if (output_at + copy_len <= output_end) {
memcpy(output_at, output_from, copy_byte);
output_at += copy_len;
} else {
copy_len >>= 1;
copy_byte >>= 1;
}
} //while
} //if
[&](ptrdiff_t j_start, ptrdiff_t j_end) {
for (auto j = j_start; j < j_end; j++) {
auto output_offset = output_offsets[j];
if (output_offset % output_dim_group[i] == 0) {
auto copy_len = output_dim_group[i] / expand_dim_size[i];
auto copy_byte = copy_len * sizeof(T);
auto output_from = output_data + output_offset;
auto output_at = output_from + copy_len;
auto output_end = output_from + output_dim_group[i];
while (output_at + copy_len <= output_end) {
memcpy(output_at, output_from, copy_byte);
output_at += copy_len;
copy_len <<= 1;
copy_byte <<= 1;
} //while
while (output_at < output_end) {
if (output_at + copy_len <= output_end) {
memcpy(output_at, output_from, copy_byte);
output_at += copy_len;
} else {
copy_len >>= 1;
copy_byte >>= 1;
}
} //while
} //if
} // for
}; //copy_fn
if (per_thread_tasks > 20) {
concurrency::ThreadPool::TrySimpleParallelFor(
concurrency::ThreadPool::TryParallelFor(
context->GetOperatorThreadPool(),
distribute_count,
copy_fn);
static_cast<double>(copy_byte),
std::move(copy_fn));
} else {
for (int64_t j = 0; j < distribute_count; ++j) {
copy_fn(j);
}
copy_fn(0, distribute_count);
} //else
} //for
return Status::OK();