未验证 提交 e5b0d9e1 编写于 作者: L liuyuhui 提交者: GitHub

[Kunlun] Add condition_variable and notify() in BindThreadedSSAGraphExecutor (#30586)

上级 ca338214
...@@ -30,9 +30,6 @@ namespace paddle { ...@@ -30,9 +30,6 @@ namespace paddle {
namespace framework { namespace framework {
namespace details { namespace details {
static std::atomic<unsigned int> exec_op_count_;
static std::atomic<int> error_state;
BindThreadedSSAGraphExecutor::BindThreadedSSAGraphExecutor( BindThreadedSSAGraphExecutor::BindThreadedSSAGraphExecutor(
const ExecutionStrategy &strategy, const std::vector<Scope *> &local_scopes, const ExecutionStrategy &strategy, const std::vector<Scope *> &local_scopes,
const std::vector<Scope *> &local_exec_scopes, const std::vector<Scope *> &local_exec_scopes,
...@@ -125,7 +122,7 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream( ...@@ -125,7 +122,7 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream(
for (auto cur_op : ready_fetch_ops) { for (auto cur_op : ready_fetch_ops) {
ready_ops->Push(cur_op); ready_ops->Push(cur_op);
} }
// Atomic variable, no need to lock
exec_op_count_ = 0; exec_op_count_ = 0;
platform::XPUPlace cur_place; platform::XPUPlace cur_place;
...@@ -134,9 +131,8 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream( ...@@ -134,9 +131,8 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream(
while (cur_count < op_deps_.size()) { while (cur_count < op_deps_.size()) {
cur_count++; cur_count++;
auto cur_op = ready_ops->Pop(); auto cur_op = ready_ops->Pop();
// when execption, get cur_op == nullptr
if (cur_op == nullptr) { if (cur_op == nullptr) {
// sleep a while to make sure worker thread quit
sleep(10);
exec_op_count_ = op_deps_.size(); exec_op_count_ = op_deps_.size();
break; break;
} }
...@@ -151,14 +147,16 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream( ...@@ -151,14 +147,16 @@ FetchResultType BindThreadedSSAGraphExecutor::RunMainStream(
RunOpAsyncMainStream(cur_op, op_deps.get(), ready_ops, cur_index); RunOpAsyncMainStream(cur_op, op_deps.get(), ready_ops, cur_index);
} }
} }
while (exec_op_count_ < op_deps_.size()) { {
std::unique_lock<std::mutex> lock(mutex_);
cv_.wait(lock, [&] { return exec_op_count_ >= op_deps_.size(); });
} }
// Wait FetchOps.
ClearFetchOp(graph_, &fetch_ops);
if (exception_.IsCaught()) { if (exception_.IsCaught()) {
ExecutionFinal(&fetch_ops); ExecutionFinal(&fetch_ops);
} }
// Wait FetchOps.
ClearFetchOp(graph_, &fetch_ops);
return fetches; return fetches;
} }
...@@ -222,7 +220,8 @@ void BindThreadedSSAGraphExecutor::InsertFetchOps( ...@@ -222,7 +220,8 @@ void BindThreadedSSAGraphExecutor::InsertFetchOps(
} }
} }
} }
// RunMultiDeviceOpAsync function is used for Communicated OPs
// like all_reduce\broadcast among multicards.
void BindThreadedSSAGraphExecutor::RunMultiDeviceOpAsync( void BindThreadedSSAGraphExecutor::RunMultiDeviceOpAsync(
OpHandleBase *op, OpHandleBase *op,
std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps, std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps,
...@@ -256,10 +255,12 @@ void BindThreadedSSAGraphExecutor::RunMultiDeviceOpAsync( ...@@ -256,10 +255,12 @@ void BindThreadedSSAGraphExecutor::RunMultiDeviceOpAsync(
ready_ops->Push(nullptr); ready_ops->Push(nullptr);
exception_.Catch(std::current_exception()); exception_.Catch(std::current_exception());
} }
// Atomic variable, no need to lock
exec_op_count_++; exec_op_count_++;
cv_.notify_all();
}); });
} }
// RunOpAsyncMainStream function is used for computed OPs
void BindThreadedSSAGraphExecutor::RunOpAsyncMainStream( void BindThreadedSSAGraphExecutor::RunOpAsyncMainStream(
OpHandleBase *op, OpHandleBase *op,
std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps, std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps,
...@@ -285,7 +286,9 @@ void BindThreadedSSAGraphExecutor::RunOpAsyncMainStream( ...@@ -285,7 +286,9 @@ void BindThreadedSSAGraphExecutor::RunOpAsyncMainStream(
ready_ops->Push(nullptr); ready_ops->Push(nullptr);
exception_.Catch(std::current_exception()); exception_.Catch(std::current_exception());
} }
// Atomic variable, no need to lock
exec_op_count_++; exec_op_count_++;
cv_.notify_all();
}); });
} }
......
...@@ -14,7 +14,9 @@ ...@@ -14,7 +14,9 @@
#pragma once #pragma once
#include <ThreadPool.h> #include <ThreadPool.h>
#include <condition_variable> // NOLINT
#include <memory> #include <memory>
#include <mutex> // NOLINT
#include <string> #include <string>
#include <unordered_map> #include <unordered_map>
#include <vector> #include <vector>
...@@ -76,6 +78,11 @@ class BindThreadedSSAGraphExecutor : public SSAGraphExecutor { ...@@ -76,6 +78,11 @@ class BindThreadedSSAGraphExecutor : public SSAGraphExecutor {
::ThreadPool prepare_pool_; ::ThreadPool prepare_pool_;
::ThreadPool multi_device_op_pool_; ::ThreadPool multi_device_op_pool_;
std::mutex mutex_;
std::condition_variable cv_;
std::atomic<unsigned int> exec_op_count_;
std::atomic<int> error_state;
void RunOpAsyncMainStream( void RunOpAsyncMainStream(
OpHandleBase *op, OpHandleBase *op,
std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps, std::unordered_map<OpHandleBase *, struct RunningItem> *op_deps,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册