未验证 提交 7e60cc63 编写于 作者: Z Zeng Jinle 提交者: GitHub

refine case when thread_num = 1 (#36201)

上级 092d45c3
...@@ -47,7 +47,16 @@ FastThreadedSSAGraphExecutor::FastThreadedSSAGraphExecutor( ...@@ -47,7 +47,16 @@ FastThreadedSSAGraphExecutor::FastThreadedSSAGraphExecutor(
<< "Change thread number to 1 because the toposort order is unique"; << "Change thread number to 1 because the toposort order is unique";
strategy_.num_threads_ = 1; strategy_.num_threads_ = 1;
} }
if (strategy_.num_threads_ > 1) {
pool_.reset(new ::ThreadPool(strategy.num_threads_)); pool_.reset(new ::ThreadPool(strategy.num_threads_));
} else {
auto nodes = ir::TopologySortOperations(*graph_);
traced_ops_.clear();
traced_ops_.reserve(nodes.size());
for (auto *node : nodes) {
traced_ops_.push_back(&node->Wrapper<OpHandleBase>());
}
}
for (auto &op : ir::FilterByNodeWrapper<OpHandleBase>(*graph_)) { for (auto &op : ir::FilterByNodeWrapper<OpHandleBase>(*graph_)) {
int dep = static_cast<int>(op->NotReadyInputSize()); int dep = static_cast<int>(op->NotReadyInputSize());
op_deps_.emplace(op, dep); op_deps_.emplace(op, dep);
...@@ -228,7 +237,7 @@ void FastThreadedSSAGraphExecutor::RunOpAsync( ...@@ -228,7 +237,7 @@ void FastThreadedSSAGraphExecutor::RunOpAsync(
OpHandleBase *op, OpHandleBase *op,
const std::shared_ptr<BlockingQueue<size_t>> &complete_q) { const std::shared_ptr<BlockingQueue<size_t>> &complete_q) {
++remaining_; ++remaining_;
this->pool_->enqueue([=] { auto func = [=] {
std::deque<OpHandleBase *> op_queue; std::deque<OpHandleBase *> op_queue;
op_queue.push_front(op); op_queue.push_front(op);
...@@ -287,7 +296,12 @@ void FastThreadedSSAGraphExecutor::RunOpAsync( ...@@ -287,7 +296,12 @@ void FastThreadedSSAGraphExecutor::RunOpAsync(
} }
--remaining_; --remaining_;
complete_q->Push(complete); complete_q->Push(complete);
}); };
if (pool_) {
pool_->enqueue(func);
} else {
func();
}
} }
void FastThreadedSSAGraphExecutor::PrepareAtomicOpDeps() { void FastThreadedSSAGraphExecutor::PrepareAtomicOpDeps() {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册