wehub-resource-sync
114 lines
3.4 KiB
C++
114 lines
3.4 KiB
C++
// Copyright (c) 2022 CINN Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/cinn/utils/multi_threading.h"
|
|
|
|
#include <glog/logging.h>
|
|
|
|
#include <future>
|
|
#include <thread>
|
|
#include <utility>
|
|
#include <vector>
|
|
#include "paddle/cinn/utils/string.h"
|
|
#include "paddle/common/enforce.h"
|
|
|
|
namespace cinn {
|
|
namespace utils {
|
|
|
|
SequenceDispatcher::SequenceDispatcher(int begin, int end, int step)
|
|
: end_(end), step_(step), index_(begin) {
|
|
PADDLE_ENFORCE_LE(
|
|
begin,
|
|
end,
|
|
::common::errors::InvalidArgument("begin[%d] > end[%d]", begin, end));
|
|
PADDLE_ENFORCE_GT(
|
|
step, 0, ::common::errors::InvalidArgument("step is less than 0."));
|
|
}
|
|
|
|
int SequenceDispatcher::Next() const {
|
|
int idx = -1;
|
|
if (index_ >= end_ || (idx = index_.fetch_add(step_)) >= end_) {
|
|
return -1;
|
|
}
|
|
|
|
return idx;
|
|
}
|
|
|
|
void parallel_run(const WorkerFuncType& fn,
|
|
JobDispatcher&& dispatcher,
|
|
int num_threads) {
|
|
if (num_threads == -1 || num_threads > std::thread::hardware_concurrency()) {
|
|
num_threads = std::thread::hardware_concurrency();
|
|
}
|
|
PADDLE_ENFORCE_GT(num_threads,
|
|
0,
|
|
::common::errors::PreconditionNotMet(
|
|
"num_threads should be greater than 0"));
|
|
|
|
// worker function of a thread
|
|
auto worker = [&fn, &dispatcher](int tid) -> int {
|
|
int index = -1, counter = 0;
|
|
while ((index = dispatcher.Next()) != -1) {
|
|
VLOG(5) << "Thread-" << tid << " process at index: " << index;
|
|
fn(index);
|
|
++counter;
|
|
}
|
|
return counter;
|
|
};
|
|
|
|
std::vector<std::future<int>> futures;
|
|
std::vector<std::thread> threads;
|
|
// The first thread runs inplace, and other `num_threads - 1` threads launched
|
|
// with std::future to run asynchronously
|
|
if (num_threads > 1) {
|
|
futures.reserve(num_threads - 1);
|
|
threads.reserve(num_threads - 1);
|
|
for (int tid = 1; tid < num_threads; ++tid) {
|
|
std::packaged_task<int(int)> task(worker);
|
|
futures.emplace_back(task.get_future());
|
|
threads.emplace_back(std::move(task), tid);
|
|
}
|
|
}
|
|
|
|
// wait results and catch exceptions
|
|
try {
|
|
int tid = 0;
|
|
int counter = worker(tid);
|
|
VLOG(4) << "Thread-0 process " << counter << " tasks.";
|
|
|
|
for (auto& future : futures) {
|
|
counter = future.get();
|
|
++tid;
|
|
VLOG(4) << "Thread-" << tid << " process " << counter << " tasks.";
|
|
}
|
|
|
|
// join threads
|
|
for (auto&& thread : threads) {
|
|
thread.join();
|
|
}
|
|
} catch (::common::EnforceNotMet& ex) {
|
|
LOG(ERROR) << ex.error_str();
|
|
PADDLE_THROW(
|
|
::common::errors::Fatal("Parallel compile Paddle enforce error"));
|
|
} catch (const std::exception& e) {
|
|
LOG(ERROR) << "Parallel compile error " << e.what();
|
|
PADDLE_THROW(::common::errors::Fatal("Parallel compile std::exception"));
|
|
} catch (...) {
|
|
PADDLE_THROW(::common::errors::Fatal("Parallel compile unknown exception"));
|
|
}
|
|
}
|
|
|
|
} // namespace utils
|
|
} // namespace cinn
|