diff --git a/onnxruntime/core/common/threadpool.cc b/onnxruntime/core/common/threadpool.cc index 3ad7ffa768..584fb3708e 100644 --- a/onnxruntime/core/common/threadpool.cc +++ b/onnxruntime/core/common/threadpool.cc @@ -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(d_of_p); + } return num_shards; } @@ -181,7 +196,7 @@ void ThreadPool::ParallelForFixedBlockSizeScheduling(const std::ptrdiff_t total, int num_work_items = static_cast(std::min(static_cast(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 run_work = [&](unsigned idx) { unsigned my_home_shard = lc.GetHomeShard(idx); unsigned my_shard = my_home_shard; diff --git a/onnxruntime/core/providers/cpu/tensor/expand.cc b/onnxruntime/core/providers/cpu/tensor/expand.cc index 4b9c3870c8..91afefca0a 100644 --- a/onnxruntime/core/providers/cpu/tensor/expand.cc +++ b/onnxruntime/core/providers/cpu/tensor/expand.cc @@ -115,68 +115,70 @@ Status Expand::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(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(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();